tensorflow2.2_实现无高级封装——以训练MNIST为例

import tensorflow as tf

# 载入MNSIT数据集
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
# ont-hot编码
y_train = tf.keras.utils.to_categorical(y_train, num_classes=10)
y_test = tf.keras.utils.to_categorical(y_test, num_classes=10)

# 定义数据集
train_data = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_data = train_data.repeat(1)
train_data = train_data.batch(32)

test_data = tf.data.Dataset.from_tensor_slices((x_test, y_test))
test_data = test_data.repeat(1)
test_data = test_data.batch(32)

# 定义模型
model = tf.keras.models.Sequential()

model.add(tf.keras.layers.Flatten(input_shape=(28,28)))
model.add(tf.keras.layers.Dense(10, activation='softmax'))

# 定义优化器
optimizer = tf.keras.optimizers.SGD(0.1)

# 计算损失和准确率
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_acc = tf.keras.metrics.CategoricalAccuracy(name='train_acc')

test_loss = tf.keras.metrics.Mean(name='test_loss')
test_acc = tf.keras.metrics.CategoricalAccuracy(name='test_acc')

# @tf.function装饰器来将python代码转成tensorflow的图表示代码，用于加速代码运行速度
# 定义训练函数
@tf.function
def train_step(data, label):
    with tf.GradientTape() as tape:
        predictions = model(data)
        loss = tf.keras.losses.MSE(label, predictions)
        # 传入loss和权值
        grad = tape.gradient(loss, model.trainable_variables)
        # 对权值进行调整
        optimizer.apply_gradients(zip(grad, model.trainable_variables))
        # 计算平均loss和准确率
        train_loss(loss)
        train_acc(label, predictions)

# 定义测试函数
@tf.function
def test_step(data, label):
    predictions = model(data)
    loss = tf.keras.losses.MSE(label, predictions)
    test_loss(loss)
    test_acc(label, predictions)

# 开始训练
for epoch in range(10):
    for img, label in train_data:
        train_step(img, label)
    for test_img, test_label in test_data:
        test_step(test_img, test_label)

    print('Epoch:%d, Loss:%.3f, Acc:%.3f, Test Loss:%.3f, Test Acc:%.3f'%(epoch, train_loss.result(), train_acc.result(), test_loss.result(), test_acc.result()))

训练结果如下：