《TensorFlow深度学习》学习笔记--10.卷积神经网络--1.LeNet-5实战

1. LeNet-5 实战

1. 导入相关的包

import tensorflow as tf
from tensorflow.keras import Sequential, layers, losses, optimizers

设置GPU（如果使用GPU，则需添加下面代码，如果不使用，则无需添加）

devices = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(devices[0], True)

2. 加载数据集并进行预处理

# 读取MNIST数据集
(x, y), (x_test, y_test) = tf.keras.datasets.mnist.load_data()	
# x.shape:(60000, 28, 28)  
# y.shape:(60000,)  
# x_test.shape:(10000, 28, 28)
# y_test.shape:(10000,)

# 自定义预处理函数
def pre_process(x, y):
    x = tf.cast(x, dtype=tf.float32) / 255.
    y = tf.cast(y, dtype=tf.int32)
    y = tf.one_hot(y, depth=10)     # one_hot 编码
    return x, y

# 构造训练集，打乱， 设置批大小， 预处理函数以及重复次数
# https://cloud.tencent.com/developer/article/1348052（参考讲解）
train_db = tf.data.Dataset.from_tensor_slices((x, y))
train_db = train_db.shuffle(buffer_size = 10000)
train_db = train_db.batch(batch_size = 512)
train_db = train_db.map(pre_process)

3. 创建网络层

# 定义 LeNet-5
network = Sequential([								# 网络容器
    layers.Conv2D(6, kernel_size=3, strides=1),		# 卷积层，6个3*3卷积核，步长：1
    layers.MaxPooling2D(pool_size=2, strides=2),	# 池化层，高宽各减半
    layers.ReLU(),									# 激活函数
    layers.Conv2D(16, kernel_size=3, strides=1),	# 卷积层，16个3*3卷积核，步长：1
    layers.MaxPooling2D(pool_size=2, strides=2),	# 池化层，高宽各减半
    layers.ReLU(),									# 激活函数
    layers.Flatten(),								# 打平层，方便全连接层处理
    layers.Dense(120, activation='relu'),			# 全连接层，120个节点
    layers.Dense(84, activation='relu'),			# 全连接层，84个节点
    layers.Dense(10)								# 全连接层，10个节点
])

network.build(input_shape=(None, 28, 28, 1))		# 构建网络模型，给定输入 X 的形状
network.summary()									# 统计网络信息

4. 模型训练

criteon = losses.CategoricalCrossentropy(from_logits=True)
optimizer = optimizers.RMSprop(0.001)
for i in range(20):                         # 训练 20 个 epoch
    for step, (x, y) in enumerate(train_db):
        with tf.GradientTape() as tape:
            x = tf.expand_dims(x, axis=3)   # [batch_size, 28, 28] -> [batch_size, 28, 28, 1]
            out = network(x)                # 前向计算 -> [batch_size, 10]
            # 在数据集预处理时已经 one_hot,所以此处不用再进行one_hot处理。
            loss = criteon(y, out)          # 计算交叉熵损失函数，标量

        grads = tape.gradient(loss, network.trainable_variables)			# 自动计算梯度
        optimizer.apply_gradients(zip(grads, network.trainable_variables))	# 自动更新参数

        if step % 100 == 0:
            print(f'epoch: {i + 1},\t step: {step} \t loss: {loss}')

network.save('LeNet-5.h5')					# 保存网络模型

5. 模型测试

# 构建测试集
test_db = tf.data.Dataset.from_tensor_slices((x_test, y_test))
test_db = test_db.shuffle(buffer_size=10000).batch(batch_size=512).map(pre_process)

# 加载训练的网络模型
network = tf.keras.models.load_model('LeNet-5.h5')
network.summary()

# 模型测试
correct, total = 0,0
for x,y in test_db: 					# 遍历所有训练集样本
    x = tf.expand_dims(x,axis=3)        # [batch_size, 28, 28] -> [batch_size, 28, 28, 1]
    out = network(x)                    # 前向计算 -> [batch_size, 10]
    pred = tf.argmax(out, axis=-1)
    y = tf.cast(y, tf.int64)
    y = y = tf.argmax(y, axis=-1)
    # 统计预测正确数量
    correct += float(tf.reduce_sum(tf.cast(tf.equal(pred, y),tf.float32)))
    total += x.shape[0]					# 统计预测样本总数
    
# 计算准确率
print('test acc:', correct/total)