tensorflow2.0-python实现Mnist分类问题-注释超详细

import tensorflow as tf
from tensorflow.keras import layers, optimizers, datasets
from tensorflow import keras
import matplotlib.pyplot as plt

# 加载minst数据集,load_data返回两个元组数据：训练集和测试集，每个元组的第一个元素是图片数据，第二个是对应的类别数字
(x, y), (x_val, y_val) = datasets.mnist.load_data()
# 转换为浮点张量，并缩放的-1~1，minst数据集数值范围是[0, 255]，通过归一化到[0, 1.]中间，再缩放到[-1, 1]之间
x = 2 * tf.convert_to_tensor(x, dtype=tf.float32) / 255. - 1
# 转换为整型张量
y = tf.convert_to_tensor(y, dtype=tf.int32)
# One-Hot编码,，指定类别总数为10
y = tf.one_hot(y, depth=10)

print(x.shape, y.shape)
# print(x.shape, y.shape)
# 构建数据集对象
train_dataset = tf.data.Dataset.from_tensor_slices((x, y))
# 批量训练
train_dataset = train_dataset.batch(200)
# 创建第一层网络，设置输出节点数为256，激活函数为ReLU
# layers.Dense(256, activation='relu')
# 利用Sequential容器分装三个网络层，前网络层的输出默认作为下一层的输入
model = keras.Sequential([
    # 3个非线性层的嵌套模型
    # 隐藏层1
    layers.Dense(512, activation='relu'),
    # 隐藏层2
    layers.Dense(256, activation='relu'),
    # 输出层，输出节点为10
    layers.Dense(10)])
optimizer = optimizers.SGD(learning_rate=0.001)


# 定义训练方法
def train():

    for epoch in range(30):
        train_epoch(epoch)
# 定义训练世代
def train_epoch(epoch):

    # 第四步 循环
    for step, (x, y) in enumerate(train_dataset):
        # 构建梯度记录环境
        with tf.GradientTape() as tape:
            # [b, 28, 28] => [b, 784]
            # -1指的是不知道有几张图片，用-1代替，有几张都这么操作
            x = tf.reshape(x, (-1, 28 * 28))
            # 第一步 得到模型输出 [b, 784] => [b, 10]
            out = model(x)
            # 第二步 计算差的平方和， [b, 10]
            loss = tf.square(out - y)
            # 计算每个样本的平均误差
            loss = tf.reduce_sum(loss) / x.shape[0]
         # 第三步 计算参数的梯度w1,w2,w3,b1,b2,b3
        # 自动计算梯度
        grads = tape.gradient(loss, model.trainable_variables)
        # w' = w - lr * grad，更新网络参数
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
        # 输出记录
        if step % 100 == 0:
            print(epoch, step, 'loss:', loss.numpy())

    return mse


if __name__ == '__main__':
    train()