1. 结构分析
LeNet网络论文： Gradient-Based Learning Applied to Document Recognition
下面用PyTorch搭建一下LeNet网络（如下图），可以看见LeNet包含一个卷积层、一个池化下采样层、一个卷积层、一个池化下采样层、3个全连接层。
注： 由于上图结构输入的图像是灰度图像，通道只有1，而本次我们使用的数据集为RGB图像，通道为3，所以每层的卷积核个数不与图中一一对应，只保持结构大致相同。
卷积层conv1: Kernel: 16 kernel_size: 5 padding: 0 stride: 1 input_size:[3,32,32] out_size: [16,28,28]
池化层pool1: kernel_size: 2 padding: 0 stride: 2 input_size: [16,28,28] out_size: [16,14,14]
卷积层conv2: Kernel: 32 kernel_size: 5 padding: 0 stride: 1 input_size: [16,14,14] out_size: [32,10,10]
池化层pool2: kernel_size: 2 padding: 0 stride: 2 input_size: [32,10,10] out_size: [32,5,5]
全连接层FC1: 3255 输出结点数：120
全连接层FC2: 输入结点数：120 输出结点数：84
全连接层FC3：输入结点数：84 输出节点数（分类数目）：10
2. 数据处理
本次训练数据集来自CIFAR10，训练集有50000张图片，测试集10000张图片，有10个类别，是一个RGB的图像数据集，在PyTorch Tensor中，维度的排列顺序为 [bath, channel, height, width]，其中bath表示一次处理图片的数量（由于图片过多，内存放不下，所以一般分批次使用训练集）。
本次数据处理比较简单主要步骤：
- 定义两个转换函数，一个用于将图片转换成tensor，另外一个将tensor进行标准化
- 载入测试数据和训练数据
- 将测试数据和训练数据打包，设置一个batch包含图片的数量，训练集将打乱顺序，而测试集按顺序打包。
from torchvision import transforms, datasets
transform = transforms.Compose(
[
transforms.ToTensor(), # 图片转tensor
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # 输入数据标准化
]
)
# 函数参数依次为：数据路径，是否为训练集，转换函数，是否下载
train_data = datasets.CIFAR10(root="../Dataset/CIFAR10", train=True, transform=transform, download=False) # 训练50000张
test_data = datasets.CIFAR10(root="../Dataset/CIFAR10", train=False, transform=transform, download=False) # 测试10000张
train_loader = DataLoader(train_data, batch_size=32, shuffle=True, num_workers= 0) # 将数据封装成一个个batch
test_loader = DataLoader(test_data, batch_size=10000, shuffle=False, num_workers=0)
3. 网络搭建
import torch.nn.functional as F
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.model = nn.Sequential(
nn.Conv2d(3, 16, 5), # input_size:[3,32,32] out_size: [16,28,28]
nn.Sigmoid(),
nn.AvgPool2d(2), # input_size: [16,28,28] out_size: [16,14,14]
nn.Conv2d(16, 32, 5), # input_size: [16,14,14] out_size: [32,10,10]
nn.Sigmoid(),
nn.AvgPool2d(2), # input_size: [32,10,10] out_size: [32,5,5]
nn.Flatten(), # 矩阵展开
nn.Linear(32 * 5 * 5, 120), nn.Sigmoid(),
nn.Linear(120, 84), nn.Sigmoid(),
nn.Linear(84, 10)
)
def forward(self, x):
x = self.model(x)
return x
4. 训练模型
一般流程如下：
-
先实例化一个网络模型的对象
-
定义损失函数
-
定义优化器
-
设置epoch，即训练轮数
-
开始训练，对于每一个epoch：
-
从DataLoader中获取每个batch的数据，这是一个可迭代的对象：
- 数据放入实例化的网络中进行前向传播得到输出
- 计算损失
- 计算梯度
- 使用优化器更新参数
-
处理完所有batch后（即整个训练集训练完一次后），使用测试集进行评估准确率
-
-
处理完所有的epoch，训练结束，保存模型参数
此外还需要注意的是，如果需要用GPU运行，则需要将一下数据都放入GPU的内存中（要是用GPU运算还需要CUDA环境的支持）：
- 实例化的模型
- 训练数据和测试数据
- 损失函数
# Pytorch中默认使用CPU，可用使用.to()函数指定使用GPU运算
net = LeNet().to(device=device) # 实例化网络并将其放入CPU中
loss_fn = nn.CrossEntropyLoss() # 实例化损失函数，这里使用交叉熵损失函数
loss_fn.to(device) # 损失函数放入GPU
optimizer = optim.Adam(net.parameters(), lr=0.005) # 定义优化器，使用Adam优化器
test_iter = iter(test_loader) # DataLoader是可迭代对象，使用iter获得迭代器
# 由于之前定义测试集的batch_size 为10000，即只有一个batch为整个数据集，所以只要取一个next就能获得所有测试集数据
test_imgs, test_labels = test_iter.next()
test_imgs = test_imgs.to(device) # 测试集数据放入GPU
test_labels = test_labels.to(device)# 测试集数据放入GPU
epoch = 20 # 设置训练轮数
print(device) # 输出一下当前本机可用设备是GPU还是CPU
print("------训练开始-----")
for i in range(epoch):
tot_loss = 0.0 # 总损失值
for data in train_loader: # 遍历训练集的DataLoader
imgs, labels = data # 获得训练集一个batch的图片和标签
imgs = imgs.to(device) # 放入GPU
labels = labels.to(device) # 放入GPU
outputs = net(imgs) # 前向传播得到输出
loss = loss_fn(outputs, labels) # 使用损失函数计算输出结果和标签的误差
optimizer.zero_grad() # 反向传播之前需要先清空一下梯度，非常重要！！不然梯度会累加
loss.backward() # 反向传播
optimizer.step() # 优化器优化
tot_loss = tot_loss + loss.item() # 累加损失值
with torch.no_grad(): # 一个epoch结束后测试模型，取消梯度跟踪
outputs = net(test_imgs) # 前向传播
predict_y = torch.max(outputs, dim=1)[1] # 得到预测结果，由于输出的是每个标签的概率，所以取最大值
accurcy = (predict_y == test_labels).sum().item() / test_labels.size(0) # 计算准确率
print("epoch= %d, loss = %.3lf, accury= %.3lf" %(i+1, tot_loss, accurcy))
tot_loss = 0.0
print("-----训练结束-----")
torch.save(net.state_dict(), "./train_result/LeNet.pth") # 保存模型
5. 预测数据
预测数据流程如下：
- 载入模型
- 载入待测试的图片，转换成tensor
- 前向传播得到每个类别的概率，取最大的输出
import torchvision.transforms as transforms
from PIL import Image
from model import LeNet
transform = transforms.Compose(
[
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
]
)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
net = LeNet().to(device=device)
net.load_state_dict(torch.load("./train_result/LeNet.pth"))
im = Image.open("./predict_img/cat2.jfif")
im = transform(im)
im = im.to(device)
im = torch.unsqueeze(im, dim=0)
with torch.no_grad():
outputs = net(im)
predict = torch.max(outputs, dim=1)[1]
print(classes[int(predict)])
