TF2 Hello World

伪代码书写格式：http://anl.sjtu.edu.cn/roster/2016/Webpage/Document/02-PseudoCode.pdf

import tensorflow as tf
from tensorflow.keras.layers import Dense, Flatten, Conv2D, Dropout
from tensorflow.keras import Model
from tensorflow.keras.models import Sequential
from tensorflow.keras.datasets import mnist


class MyModel(Model):
    def __init__(self):
        super().__init__()
        self.featurizer = Sequential([
            Conv2D(32, (3, 3), 1, activation='relu'),
            Flatten(),
            Dense(128, activation='relu'),
            Dropout(0.2),
            Dense(64, activation='relu'),
            Dropout(0.2),
        ])
        self.classifier = Dense(10, activation='softmax')

    def call(self, inputs, training=None, mask=None):
        x = self.featurizer(inputs)
        return self.classifier(x)

    def get_config(self):
        pass


@tf.function
def train_step(images, labels):  # batch train
    with tf.GradientTape() as tape:
        predictions = model(images)
        loss_values = loss_function(labels, predictions)
    gradients = tape.gradient(loss_values, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    train_evaluator.batch_evaluate(loss_values, labels, predictions)


@tf.function
def test_step(images, labels):  # batch test
    predictions = model(images)
    loss_values = loss_function(labels, predictions)
    test_evaluator.batch_evaluate(loss_values, labels, predictions)


class Evaluator:
    def __init__(self, name):
        self.loss = tf.keras.metrics.Mean(name=name + "loss")
        self.accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name=name + 'accuracy')

    def reset(self):
        self.loss.reset_states()
        self.accuracy.reset_states()

    def batch_evaluate(self, loss, y_true, y_pred):
        self.loss(loss)
        self.accuracy(y_true, y_pred)


if __name__ == '__main__':
    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    x_train, x_test = x_train / 255.0, x_test / 255.0

    # Add a channels dimension
    x_train = x_train[..., tf.newaxis]
    x_test = x_test[..., tf.newaxis]

    train_ds = tf.data.Dataset.from_tensor_slices(
        (x_train, y_train)).shuffle(10000).batch(32)
    test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)

    model = MyModel()
    loss_function = tf.keras.losses.SparseCategoricalCrossentropy()
    optimizer = tf.keras.optimizers.Adam()
    train_evaluator = Evaluator(name="train")
    test_evaluator = Evaluator(name="test")

    EPOCHS = 5
    for epoch in range(EPOCHS):
        # 在下一个epoch开始时，重置评估指标
        train_evaluator.reset()
        test_evaluator.reset()

        # 训练并评估
        for train_images, train_labels in train_ds:
            train_step(train_images, train_labels)

        # 测试并评估
        for test_images, test_labels in test_ds:
            test_step(test_images, test_labels)

        # 输出epoch结果
        template = 'Epoch {}, Loss: {:.2f}, Accuracy: {:.2f}%, Test Loss: {:.2f}, Test Accuracy: {:.2f}%'
        print(template.format(epoch + 1,
                              train_evaluator.loss.result(),
                              train_evaluator.accuracy.result() * 100,
                              test_evaluator.loss.result(),
                              test_evaluator.accuracy.result() * 100))