TensorFlow-2.0 lesson06-手写数字识别初体验

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

import tensorflow as tf  # 导入 TF 库
from tensorflow import keras # 导入 TF 子库 keras
from tensorflow.keras import layers, optimizers, datasets # 导入 TF 子库等

(x, y), (x_val, y_val) = datasets.mnist.load_data()  #加载 MNIST 数据集 #60000样本
x = tf.convert_to_tensor(x, dtype=tf.float32)/255.  # 转换为浮点张量
y = tf.convert_to_tensor(y, dtype=tf.int32)  # 转换为整形张量
y = tf.one_hot(y, depth=10)  # one-hot 编码
print('x shape: ', x.shape,  'y shape: ', y.shape)

train_dataset = tf.data.Dataset.from_tensor_slices((x, y)) # 构建数据集对象
train_dataset = train_dataset.batch(200) # 批量训练，每个Batch大小有200个样本，模型更新参数一次

model = keras.Sequential([  # 3 个非线性层的嵌套模型
    layers.Dense(512, activation='relu'), # 创建一层网络，设置输出节点数为512，激活函数类型为 ReLU 隐藏层1
    layers.Dense(256, activation='relu'),# 创建一层网络，设置输出节点数为256，激活函数类型为 ReLU 隐藏层2
    layers.Dense(10)]) # 输出层，输出节点数为 10

optimizer = optimizers.SGD(learning_rate = 0.01) #优化器 learning rate

def train_epoch(epoch):
    # Step4.loop
    for step, (x, y) in enumerate(train_dataset):#每200个样本为一批，更新一次参数，总共60000样本*200=300个batch
        with tf.GradientTape() as tape: #新创建的GradientTape类对象
            x = tf.reshape(x, (-1, 28*28))#打平
            # Step1. compute output
            # [b, 784] => [b, 10]
            out = model(x)#输入网络模型
            # Step2. compute loss
            loss = tf.reduce_sum(tf.square(out - y)) / x.shape[0]
        # Step3. optimize and update w1, w2, w3, b1, b2, b3
        # 根据tape上面的上下文来计算某个或者某些tensor的梯度
        grads = tape.gradient(loss, model.trainable_variables)
        # w' = w - lr * grad 把计算出来的梯度更新到变量上面去。
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

        print('epoch:',epoch, 'step:', step, 'loss:', loss.numpy())

def train():
    for epoch in range(60):#有60个epoch，模型将暴露或传递整个数据集60次，共60*300个batch
        train_epoch(epoch)

if __name__ == '__main__':
    train()

参考
tf.GradientTape类相关说明

优化器  optimizers