TensorFlow2 模型建立与训练

首先，强力推荐大神关于Tensoflow2的介绍文档！文档链接

模型的构建： tf.keras.Model 和 tf.keras.layers
模型的损失函数： tf.keras.losses
模型的优化器： tf.keras.optimizer
模型的评估： tf.keras.metrics

任务介绍：自定义模型完成 MNIST 手写体数字图片数据集的分类任务！
代码如下：

import tensorflow as tf
import numpy as np

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=10)

    def __call__(self, inputs):
        x = self.flatten(inputs)
        x = self.dense1(x)
        x = self.dense2(x)
        output = tf.nn.softmax(x)
        return output

class MNISTLoader():
    def __init__(self):
        mnist = tf.keras.datasets.mnist
        
        #方式1：从https://s3.amazonaws.com/img-datasets/mnist.npz 下载数据，放在用户目录下的.keras/dataset下，
        #Windows 下用户目录为 C:\Users\用户名 ，Linux 下用户目录为 /home/用户名 
        (self.trainData, self.trainLabel), (self.testData, self.testLabel) = mnist.load_data(path='mnist.npz')
        #方式2：直接从官网在线加载数据
        #(self.trainData, self.trainLabel), (self.testData, self.testLabel) = mnist.load_data(path='mnist.npz')
        
        #将数据从0-255转到0-1，并且增加一个颜色通道维度到最后一维
        self.trainData = tf.expand_dims(self.trainData.astype(np.float32) / 255, -1)
        self.testData = tf.expand_dims(self.testData.astype(np.float32) / 255, -1)
        #将label转化为整数
        self.trainLabel = self.trainLabel.astype(np.int32)
        self.testLabel = self.testLabel.astype(np.int32)
        self.trainSize, self.testSize = self.trainData.shape[0], self.testData.shape[0]

    def get_batch(self, batchSize):
        #随机采样batchSize个数据返回
        index = np.random.randint(0, self.trainData.shape[0], batchSize).astype(np.int32)
        #print(self.trainData[[1,2,3],1])
        return tf.gather(self.trainData, axis=0, indices=index), tf.gather(self.trainLabel, axis=0, indices=index)

#自定义评估指标需要继承 tf.keras.metrics.Metric 类，并重写 __init__ 、 update_state 和 result 三个方法。
class SparseCategoricalAccuracy(tf.keras.metrics.Metric):
    def __init__(self):
        super().__init__()
        self.total = self.add_weight(name='total', dtype=tf.int32, initializer=tf.zeros_initializer())
        self.count = self.add_weight(name='count', dtype=tf.int32, initializer=tf.zeros_initializer())

    def update_state(self, y_true, y_pred, sample_weight=None):
        values = tf.cast(tf.equal(y_true, tf.argmax(y_pred, axis=-1, output_type=tf.int32)), tf.int32)
        self.total.assign_add(tf.shape(y_true)[0])
        self.count.assign_add(tf.reduce_sum(values))

    def result(self):
        return self.count / self.total

dataLoader = MNISTLoader()
num_epochs = 5
batch_size = 50
learning_rate = 0.001

batchNum = int(dataLoader.trainSize // batch_size * num_epochs)
model = MLP()
optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
for i in range(batchNum):
    trainData, trainLabel = dataLoader.get_batch(batch_size)
    with tf.GradientTape() as tape:
        lable_pre = model(trainData)
        loss = tf.losses.sparse_categorical_crossentropy(y_true=trainLabel, y_pred=lable_pre)
        loss = tf.reduce_mean(loss)
        print("batch %d, loss %f" % (i, loss))
    gradient = tape.gradient(loss, model.variables)
    optimizer.apply_gradients(grads_and_vars=zip(gradient, model.variables))

#定义准确率计算工具
sparse_categorical_accuracy = SparseCategoricalAccuracy()
batchNumTest = int(dataLoader.testSize // batch_size)
for i in range(batchNumTest):
    starIndex, endIndex = i*batch_size, (i+1)*batch_size
    testData = dataLoader.testData[starIndex:endIndex]
    testPre = model.predict(testData)
    #print(testPre.shape)
    sparse_categorical_accuracy.update_state(y_true=dataLoader.testLabel[starIndex:endIndex], y_pred=testPre)

print("test accuracy %f" %(sparse_categorical_accuracy.result()))