Tensorflow2 — MNIST

%matplotlib widget
import matplotlib.pyplot as plt
import tensorflow as tf
import numpy as np

# check gpu
tf.config.list_physical_devices('GPU') 
tf.test.is_built_with_cuda()
## output True

class MLP(tf.keras.Model):
    
    def __init__(self):
        super().__init__()
        self.flatten = tf.keras.layers.Flatten()
        self.dense1 = tf.keras.layers.Dense(units=100, activation=tf.nn.relu)
        self.dense2 = tf.keras.layers.Dense(units=20, activation=tf.nn.leaky_relu)
        self.dense3 = tf.keras.layers.Dense(units=10)

    @tf.function
    def call(self, inputs):  # [batch_size, 28, 28, 1]
        flat1 = self.flatten(inputs)  # [batch_size, 784]
        dens1 = self.dense1(flat1)  # [batch_size, 100]
        dens2 = self.dense2(dens1)  # [batch_size, 20]
        dens3 = self.dense3(dens2)  # [batch_size, 10]
        output = tf.nn.softmax(dens3)
        return output

# @tf.function
def one_batch_step(X, y, **kwargs):
    with tf.GradientTape() as tape:
        y_pred = model(X)
        loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred)
        loss = tf.reduce_mean(loss)
        tf.print(f"{batch_index} loss {loss}", [loss])
        with summary_writer.as_default():
            tf.summary.scalar("loss", loss, step=batch_index)
    grads = tape.gradient(loss, model.variables)
    optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))
    
for epoch_index in range(num_epochs):
    for batch_index in range(num_batches):
        X, y = data_loader.get_batch(batch_size)
        one_batch_step(X, y, batch_index=batch_index)
with summary_writer.as_default():
    tf.summary.trace_export(name="model_trace", step=0, profiler_outdir=log_dir)
    
tf.saved_model.save(model, f"saved/{model_name}")

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(100, activation=tf.nn.relu),
    tf.keras.layers.Dense(units=20, activation=tf.nn.leaky_relu),
    tf.keras.layers.Dense(10),
    tf.keras.layers.Softmax()
])


model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
    loss=tf.keras.losses.sparse_categorical_crossentropy,
    metrics=[tf.keras.metrics.sparse_categorical_accuracy]
)

model.fit(data_loader.train_data, data_loader.train_label, epochs=num_epochs, batch_size=batch_size)

print(model.evaluate(data_loader.test_data, data_loader.test_label))

model_name = "mnist_model2"
tf.saved_model.save(model, f"saved/{model_name}")