TensorFlow2.0之LeNet-5实战
from tensorflow.keras import Sequential, layers, losses, optimizers, datasets
import tensorflow as tf
def main():
# 网络容器
network = Sequential([
# 第一个卷积层,6个3*3卷积核
layers.Conv2D(6, kernel_size=3, strides=1),
# 高宽各减半的池化层
layers.MaxPooling2D(pool_size=2, strides=2),
# 激活函数
layers.ReLU(),
# 第二个卷积层,16个3*3卷积核
layers.Conv2D(16, kernel_size=3, strides=1),
# 高宽各减半的池化层
layers.MaxPooling2D(pool_size=2, strides=2),
# 激活函数
layers.ReLU(),
# 打平层,方便全连接处理
layers.Flatten(),
# 全连接层,120个节点
layers.Dense(120, activation='relu'),
# 全连接层,84节点
layers.Dense(84, activation='relu'),
# 全连接层,10个节点
layers.Dense(10)
])
# build一次网络模型,给输入x的形状,其中4为随意给的batchsize
network.build(input_shape=(4, 28, 28, 1))
# 统计网络信息
network.summary()
# 创建损失函数的类,在实际计算时直接调用类实例即可
criteon = losses.CategoricalCrossentropy(from_logits=True)
# 构建梯度记录环境
with tf.GradientTape() as tape:
# 插入通道维度 [b,28,28] => [b,28,28,1]
x = tf.expand_dims(x, axis=3)
# 前向计算,获得10类别的概率分布,[b,784] => [b,10]
out = network(x)
# 真实标签one-hot编码,[b] => [b,10]
y_onehot = tf.one_hot(y, depth=10)
# 计算交叉熵损失函数,标量
loss = criteon(y_onehot, out)
# 自动计算梯度
grads = tape.gradient(loss, network.trainable_variables)
# 自动更新参数
optimizers.apply_gradients(zip(grads, network.trainable_variables))
if step % 50 == 0:
print(epoch, step, 'loss:', float(loss))
# 记录预测正确的数量,总样本数量
correct, total = 0, 0
# 变历所有训练集样本
for x, y in db_test:
# 插入通道维度:[b,28,28] => [b,28,28,1]
x = tf.expand_dims(x, axis=3)
# 前向计算,获得10类别的预测分布,[b,784] => [b,10]
out = network(x)
# 真实的流程时先经过Softmax再argmax
pred = tf.argmax(out, axis=-1)
y = tf.cast(y, tf.int64)
# 统计预测正确数量
correct += float(tf.reduce_sum(tf.cast(tf.equal(pred, y), tf.float32)))
# 统计预测样本总数
total += x.shape[0]
# 计算准确率
print('test acc:', correct / total)