代码已上传至github(麻烦Star~)
1.数据集介绍
利用torchvision.datasets函数可以在线导入pytorch中的数据集，包含一些常见的数据集如MNIST、CIFAR-10等。本次使用的是CIFAR10数据集，也是一个很经典的图像分类数据集，由 Hinton 的学生 Alex Krizhevsky 和 Ilya Sutskever 整理的一个用于识别普适物体的小型数据集，一共包含 10 个类别的 RGB 彩色图片。
PyTorch的CIFAR-10数据集有时下载不了，我这里将下载好的压缩包放在网盘中，需要的可以自行下载，解压后放在当前项目文件的data文件夹下。链接：https://pan.baidu.com/s/1NBHp0SxEOJ5EIyYUsDHm_g
提取码：qp3k
2.LeNet网络介绍
LeNet网络之前在我的博客详细讲解过：https://blog.csdn.net/muye_IT/article/details/123539199?spm=1001.2014.3001.5501
LeNet网络架构总览图：
3. model.py 创建
model.py ——定义LeNet网络模型
# 使用torch.nn包来构建神经网络. import torch.nn as nn import torch.nn.functional as F class LeNet(nn.Module): # 继承于nn.Module这个父类 def __init__(self): # 初始化网络结构 super(LeNet, self).__init__() # 多继承需用到super函数 self.conv1 = nn.Conv2d(3, 16, 5) self.pool1 = nn.MaxPool2d(2, 2) self.conv2 = nn.Conv2d(16, 32, 5) self.pool2 = nn.MaxPool2d(2, 2) self.fc1 = nn.Linear(32*5*5, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) def forward(self, x): # 正向传播过程 x = F.relu(self.conv1(x)) # input(3, 32, 32) output(16, 28, 28) x = self.pool1(x) # output(16, 14, 14) x = F.relu(self.conv2(x)) # output(32, 10, 10) x = self.pool2(x) # output(32, 5, 5) x = x.view(-1, 32*5*5) # output(32*5*5) x = F.relu(self.fc1(x)) # output(120) x = F.relu(self.fc2(x)) # output(84) x = self.fc3(x) # output(10) return x
Conv2d、MaxPool2d、Linear在pytorch中对应的函数，以及函数参数的设置
常见的参数：
- in_channels：输入特征矩阵的深度。如输入一张RGB彩色图像，那in_channels=3
- out_channels：输入特征矩阵的深度。也等于卷积核的个数，使用n个卷积核输出的特征矩阵深度就是n
- kernel_size：卷积核的尺寸。可以是int类型，如3 代表卷积核的height=width=3，也可以是tuple类型如(3,5)代表卷积核的height=3，width=5
- stride：卷积核的步长。默认为1，和kernel_size一样输入可以是int型，也可以是tuple类型
- padding：补零操作，默认为0。可以为int型如1即补一圈0，如果输入为tuple型如(2, 1) 代表在上下补2行，左右补1列
Conv2d ['stride', 'padding', 'dilation', 'groups','padding_mode', 'output_padding', 'in_channels','out_channels', 'kernel_size'] MaxPool2d（'kernel_size', 'stride', 'padding', 'dilation','return_indices', 'ceil_mode'） Linear('in_features', 'out_features'）
4. train.py创建
train.py ——加载数据集并训练，训练集计算loss，测试集计算accuracy，保存训练好的网络参数
4.1 相关包的加载
import torch import torchvision import torch.nn as nn from model import LeNet import torch.optim as optim import torchvision.transforms as transforms
4.2 数据预处理
由shape (H x W x C) in the range [0, 255] → shape (C x H x W) in the range [0.0, 1.0]
transform = transforms.Compose( [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
4.3 加载数据集
# 导入50000张训练图片 train_set = torchvision.datasets.CIFAR10(root='./data', # 数据集存放目录 train=True, # 表示是数据集中的训练集 download=True, # 第一次运行时为True，下载数据集，下载完成后改为False transform=transform) # 预处理过程 # 加载训练集，实际过程需要分批次（batch）训练 train_loader = torch.utils.data.DataLoader(train_set, # 导入的训练集 batch_size=50, # 每批训练的样本数 shuffle=False, # 是否打乱训练集 num_workers=0) # 使用线程数，在windows下设置为0
4.4 加载测试集
# 导入10000张测试图片 test_set = torchvision.datasets.CIFAR10(root='./data', train=False, # 表示是数据集中的测试集 download=False,transform=transform) # 加载测试集 test_loader = torch.utils.data.DataLoader(test_set, batch_size=10000, # 每批用于验证的样本数 shuffle=False, num_workers=0) # 获取测试集中的图像和标签，用于accuracy计算 test_data_iter = iter(test_loader) test_image, test_label = test_data_iter.next()
4.5 代码（GPU训练版本）
使用下面语句可以在有GPU时使用GPU，无GPU时使用CPU进行训练
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(device)
也可以直接指定
device = torch.device("cuda") # 或者 # device = torch.device("cpu")
对应的，需要用to()函数来将Tensor在CPU和GPU之间相互移动，分配到指定的device中计算
import torch import torchvision import torch.nn as nn from model import LeNet import torch.optim as optim import torchvision.transforms as transforms def main(): transform = transforms.Compose( [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) # 50000张训练图片 # 第一次使用时要将download设置为True才会自动去下载数据集 train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=False, transform=transform) train_loader = torch.utils.data.DataLoader(train_set, batch_size=36, shuffle=True, num_workers=0) # 10000张验证图片 # 第一次使用时要将download设置为True才会自动去下载数据集 val_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=False, transform=transform) val_loader = torch.utils.data.DataLoader(val_set, batch_size=5000, shuffle=False, num_workers=0) val_data_iter = iter(val_loader) val_image, val_label = val_data_iter.next() # classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') net = LeNet() net.to(device) # 将网络分配到指定的device中 loss_function = nn.CrossEntropyLoss() optimizer = optim.Adam(net.parameters(), lr=0.001) for epoch in range(5): # loop over the dataset multiple times running_loss = 0.0 for step, data in enumerate(train_loader, start=0): # get the inputs; data is a list of [inputs, labels] inputs, labels = data # zero the parameter gradients optimizer.zero_grad() # forward + backward + optimize outputs = net(inputs.to(device)) # 将inputs分配到指定的device中 loss = loss_function(outputs, labels.to(device)) # 将labels分配到指定的device中 loss.backward() optimizer.step() # print statistics running_loss += loss.item() if step % 500 == 499: # print every 500 mini-batches with torch.no_grad(): outputs = net(test_image.to(device)) # 将test_image分配到指定的device中 predict_y = torch.max(outputs, dim=1)[1] accuracy = torch.eq(predict_y, val_label).sum().item() / val_label.size(0) print('[%d, %5d] train_loss: %.3f test_accuracy: %.3f' % (epoch + 1, step + 1, running_loss / 500, accuracy)) running_loss = 0.0 print('Finished Training') save_path = './Lenet.pth' torch.save(net.state_dict(), save_path) if __name__ == '__main__': main()
4.6 代码（CPU训练版本）
import torch import torchvision import torch.nn as nn from model import LeNet import torch.optim as optim import torchvision.transforms as transforms def main(): transform = transforms.Compose( [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) # 50000张训练图片 # 第一次使用时要将download设置为True才会自动去下载数据集 train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=False, transform=transform) train_loader = torch.utils.data.DataLoader(train_set, batch_size=36, shuffle=True, num_workers=0) # 10000张验证图片 # 第一次使用时要将download设置为True才会自动去下载数据集 val_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=False, transform=transform) val_loader = torch.utils.data.DataLoader(val_set, batch_size=5000, shuffle=False, num_workers=0) val_data_iter = iter(val_loader) val_image, val_label = val_data_iter.next() # classes = ('plane', 'car', 'bird', 'cat', # 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') net = LeNet() loss_function = nn.CrossEntropyLoss() optimizer = optim.Adam(net.parameters(), lr=0.001) for epoch in range(5): # loop over the dataset multiple times running_loss = 0.0 for step, data in enumerate(train_loader, start=0): # get the inputs; data is a list of [inputs, labels] inputs, labels = data # zero the parameter gradients optimizer.zero_grad() # forward + backward + optimize outputs = net(inputs) loss = loss_function(outputs, labels) loss.backward() optimizer.step() # print statistics running_loss += loss.item() if step % 500 == 499: # print every 500 mini-batches with torch.no_grad(): outputs = net(val_image) # [batch, 10] predict_y = torch.max(outputs, dim=1)[1] accuracy = torch.eq(predict_y, val_label).sum().item() / val_label.size(0) print('[%d, %5d] train_loss: %.3f test_accuracy: %.3f' % (epoch + 1, step + 1, running_loss / 500, accuracy)) running_loss = 0.0 print('Finished Training') save_path = './Lenet.pth' torch.save(net.state_dict(), save_path) if __name__ == '__main__': main()
5. predict.py 创建
predict.py——得到训练好的网络参数后，用自己找的图像进行分类测试，自己下载一张照片保存在根目录下，命名为1.jpg
import torch import torchvision.transforms as transforms from PIL import Image from model import LeNet def main(): transform = transforms.Compose( [transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') net = LeNet() net.load_state_dict(torch.load('Lenet.pth')) im = Image.open('1.jpg')#自己下载一张照片保存在根目录下，命名为1.jpg im = transform(im) # [C, H, W] im = torch.unsqueeze(im, dim=0) # [N, C, H, W] with torch.no_grad(): outputs = net(im) predict = torch.max(outputs, dim=1)[1].data.numpy() print(classes[int(predict)]) if __name__ == '__main__': main()
