diff --git a/integration_tests/numerical_test.py b/integration_tests/numerical_test.py new file mode 100644 index 000000000..5ab52334b --- /dev/null +++ b/integration_tests/numerical_test.py @@ -0,0 +1,88 @@ +import numpy as np +from tensorflow import keras + +import keras_core + +NUM_CLASSES = 10 + + +def build_mnist_data(num_classes): + (x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data() + + # Scale images to the [0, 1] range + x_train = x_train.astype("float32") / 255 + x_test = x_test.astype("float32") / 255 + # Make sure images have shape (28, 28, 1) + x_train = np.expand_dims(x_train, -1) + x_test = np.expand_dims(x_test, -1) + + # convert class vectors to binary class matrices + y_train = keras.utils.to_categorical(y_train, num_classes) + y_test = keras.utils.to_categorical(y_test, num_classes) + + print("x_train shape:", x_train.shape) + print(x_train.shape[0], "train samples") + print(x_test.shape[0], "test samples") + + return x_train, y_train, x_test, y_test + + +def build_keras_model(keras_module, num_classes): + + input_shape = (28, 28, 1) + + model = keras_module.Sequential( + [ + keras_module.Input(shape=input_shape), + keras_module.layers.Conv2D(32, kernel_size=(3, 3), activation="relu"), + keras_module.layers.MaxPooling2D(pool_size=(2, 2)), + keras_module.layers.Conv2D(64, kernel_size=(3, 3), activation="relu"), + keras_module.layers.MaxPooling2D(pool_size=(2, 2)), + keras_module.layers.Flatten(), + keras_module.layers.Dropout(0.5), + keras_module.layers.Dense(num_classes, activation="softmax"), + ] + ) + + model.summary() + return model + + +def train_model(model, x, y): + batch_size = 128 + epochs = 1 + + model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"]) + + return model.fit(x, y, batch_size=batch_size, epochs=epochs, validation_split=0.1) + + +def eval_model(model, x, y): + score = model.evaluate(x, y, verbose=0) + print("Test loss:", score[0]) + print("Test accuracy:", score[1]) + return score + + +def numerical_test(): + x_train, y_train, x_test, y_test = build_mnist_data(NUM_CLASSES) + keras_model = build_keras_model(keras, NUM_CLASSES) + keras_core_model = build_keras_model(keras_core, NUM_CLASSES) + + # Make sure both model have same weights before training + keras_core_model.set_weights(keras_model.weights) + + for kw, kcw in zip(keras_model.weights, keras_core_model.weights): + np.testing.assert_allclose(kw, kcw) + + keras_history = train_model(keras_model, x_train, y_train) + keras_core_history = train_model(keras_core_model, x_train, y_train) + + for h, ch in zip(keras_history.history.items(), keras_core_history.history.items()): + # They are not exactly equal at the moment. + print(h) + print(ch) + + +if __name__ == "__main__": + numerical_test()