pytorch-官方demo-LeNet搭建训练测试全流程

一、搭建环境

主要搭建的是pytorch1.7.1+ python3.7 + cuda11.0环境

具体搭建过程可以参考其他优秀博客

你可能遇到的一些问题可以参考我的博客：

1.Error #15:Initializing libiomp5md.dll, but found libiomp5md.dll already initialized.

2.【torch.cuda.is_available()输出false】

3.ERROR: Could notfind a version thatsatisfies the requirement

等等

搭建过程链接

二、Lenet网络介绍

lenet网络是卷积神经网络的雏形，输入的是但通道的灰度图片，大小为32 * 32

处理过程：

输入（input）特征图尺寸	处理操作	输出（output）特征图尺寸
1 * 32 * 32	卷积层（6卷积核大小为1 * 5 * 5） + relu	6 * 28 * 28
6 *28 * 28	池化层（最大下采样 大小为2 * 2）	6 * 14 * 14
6 * 14 * 14	卷积层（16卷积核大小为6 * 5 * 5）+relu	16 * 10 * 10
16 * 10 * 10	池化层（最大下采样 大小为2 * 2）	16 * 5 * 5
16 * 5 * 5	全连接层 + relu	120
120	全连接层 + relu	84
84	全连接层 + softmax	10

在本次demo使用的是cifar10数据集，每张图片为彩色图片有三个通道

pytorch tensor通道排列顺序为[batch, channel, height, width]

batch为一批图像大小，channel为图片通道数，height核width为图片高和宽

看官方源码：Conv2d方法

我们可以复制Conv2d到pytorch的Docs搜索

得到：

地址：Conv2d — PyTorch 1.13 documentation可以查看函数具体原理

三、model.py文件

import torch.nn as nn
import torch.nn.functional as F


class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()  # 这里涉及继承，所以使用super来避免一些问题
        self.conv1 = nn.Conv2d(3, 16, 5)  # 卷积层 输入通道数为3, 卷积核个数为16， 卷积核大小为5 * 5
        self.pool1 = nn.MaxPool2d(2, 2)  # 池化层 池化大小2 * 2，步长为2
        self.conv2 = nn.Conv2d(16, 32, 5)  # 卷积层 输入通道数为16, 卷积核个数为32， 卷积核大小为5 * 5
        self.pool2 = nn.MaxPool2d(2, 2)  # 池化层 池化大小2 * 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)  # input(16, 28, 28) output(16, 14, 14)
        x = F.relu(self.conv2(x))  # input(16, 14, 14) output(32, 10, 10)
        x = self.pool2(x)  # input(32, 10, 10) output(32, 5, 5)
        x = x.view(-1, 32 * 5 * 5)
        x = F.relu(self.fc1(x))  # input(32 * 5 * 5) output(120)
        x = F.relu(self.fc2(x))  # input(120) output(84)
        x = self.fc3(x)  # input(84) output(10)
        return x

插入本次实验用的数据集CIFAR10数据集

介绍地址：Training a Classifier — PyTorch Tutorials1.13.1+cu117 documentation

四、train.py

下面为train.py中的下载数据集代码部分

import torch
import torchvision
import torch.nn as nn
from model import LeNet
import torch.optim as optim
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"


# transforms.ToTensor()转变为tensor数据将其归一化，
# 然后transforms.Normalize((0.5, 0.5, 0.5),(0.5, 0.5, 0.5))再标准化到[0, 1]
# 这里均值为0.5,标准差也为0.5,三个0.5并列是指三个通道维度都要标准化
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 50000张训练图片
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=False, transform=transform)

将download设置为true即可下载，有点慢，或者可以找其他方式下载数据集

为了查看测试集的图片，有代码如下

# 批次大小为36，shuffle表示是否打乱数据，num_workers在windows系统只能设置为0
trainloader = torch.utils.data.DataLoader(trainset, batch_size=36, shuffle=True, num_workers=0)

# 10000张测试图片
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=False, transform=transform)
# 批次大小为10000，shuffle表示是否打乱数据，num_workers在windows系统只能设置为0
testloader = torch.utils.data.DataLoader(testset, batch_size=10000, shuffle=False, num_workers=0)
# 构造迭代器
test_data_iter = iter(testloader)
test_image, test_label = test_data_iter.next()

classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')


# 展示图片函数
def imshow(img):
    img = img / 2 + 0.5  # 相等于unnormalize
    npimg = img.numpy()
    # 这里是因为之前transforms.ToTensor()转变为tensor数据将其归一化时改变了通道顺序,这里改回了[H, W, C]
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


# print labels
print(''.join('%7s' % classes[test_label[j]] for j in range(4)))
# show image
imshow(torchvision.utils.make_grid(test_image))

运行train.py（运行前记得把trainset中的download设置为False）

有可能发现了报错：

主要原因是anaconda环境存在两个libiomp5md.dll

解决方法：可以参考我的另一篇博文：链接

报错解决,输出四张图片与标签

之后记得注释这一段（展示图片函数那一行下面代码）

继续在train.py中添加代码

# 实例化模型
net = LeNet()

# 定义损失函数
loss_function = nn.CrossEntropyLoss()
# 定义优化器 第一个参数为需要训练的参数，第二个为学习率

最后一个全连接层没有使用softmax函数主要原因是，在这里调用交叉熵损失函数时已经包括了softmax函数

以下是train.py完整代码(cpu版)

import torch
import torchvision
import torch.nn as nn
from model import LeNet
import torch.optim as optim
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"


# transforms.ToTensor()转变为tensor数据将其归一化，
# 然后transforms.Normalize((0.5, 0.5, 0.5),(0.5, 0.5, 0.5))再标准化到[0, 1]
# 这里均值为0.5,标准差也为0.5,三个0.5并列是指三个通道维度都要标准化
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 50000张训练图片
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=False, transform=transform)
# 批次大小为36，shuffle表示是否打乱数据，num_workers在windows系统只能设置为0
trainloader = torch.utils.data.DataLoader(trainset, batch_size=36, shuffle=True, num_workers=0)

# 10000张测试图片
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=False, transform=transform)
# 批次大小为10000，shuffle表示是否打乱数据，num_workers在windows系统只能设置为0
testloader = torch.utils.data.DataLoader(testset, batch_size=10000, shuffle=False, num_workers=0)
# 构造迭代器
test_data_iter = iter(testloader)
test_image, test_label = test_data_iter.next()

classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')


# # 展示图片函数
# def imshow(img):
#     img = img / 2 + 0.5  # 相等于unnormalize
#     npimg = img.numpy()
#     # 这里是因为之前transforms.ToTensor()转变为tensor数据将其归一化时改变了通道顺序,这里改回了[H, W, C]
#     plt.imshow(np.transpose(npimg, (1, 2, 0)))
#     plt.show()
#
#
# # print labels
# print(''.join('%7s' % classes[test_label[j]] for j in range(4)))
# # show image
# imshow(torchvision.utils.make_grid(test_image))


# 实例化模型
net = LeNet()

# 定义损失函数
loss_function = nn.CrossEntropyLoss()
# 定义优化器 第一个参数为需要训练的参数，第二个为学习率
optimizer = optim.Adam(net.parameters(), lr=0.001)

# 训练过程
for epoch in range(10):
    running_loss = 0    # 用来垒加训练过程的损失
    for step, data in enumerate(trainloader, start=0):
        # 同时输入一个批次的的图片和标签
        inputs, labels = data
        # 下面这个方法主要用来将一个batch的导数清0，因为计算是一个批次一个批次处理的
        optimizer.zero_grad()
        # forward + backward + optimize
        output = net(inputs)
        loss = loss_function(output, labels)
        loss.backward()
        optimizer.step()    # 参数更新

        # print statistics
        running_loss += loss.item()
        if step % 500 == 499:
            with torch.no_grad():
                outputs = net(test_image)
                predict_y = torch.max(outputs, dim=1)[1]
                accuracy = (predict_y == test_label).sum().item() / test_label.size(0)
                print('[%d, %5d] train_loss:%3f test_accuracy: %3f' %
                      (epoch + 1, step + 1, running_loss/500, accuracy))
                running_loss = 0
print("finish training")
save_path = './LeNet.pth'
torch.save(net.state_dict(), save_path)

训练结果

1000是由于总共50000张图片，每个批次有36张，总共1388.88批次，也即1389步，每500步测试一次，所以会打印两次到1000

将cpu改为gpu运行版的train.py完整版代码（gpu版），其实就是将模型和数据都改为了cuda的数据格式

import torch
import torchvision
import torch.nn as nn
from model import LeNet
import torch.optim as optim
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"


# transforms.ToTensor()转变为tensor数据将其归一化，
# 然后transforms.Normalize((0.5, 0.5, 0.5),(0.5, 0.5, 0.5))再标准化到[0, 1]
# 这里均值为0.5,标准差也为0.5,三个0.5并列是指三个通道维度都要标准化
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 50000张训练图片
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=False, transform=transform)
# 批次大小为36，shuffle表示是否打乱数据，num_workers在windows系统只能设置为0
trainloader = torch.utils.data.DataLoader(trainset, batch_size=36, shuffle=True, num_workers=0)

# 10000张测试图片
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=False, transform=transform)
# 批次大小为10000，shuffle表示是否打乱数据，num_workers在windows系统只能设置为0
testloader = torch.utils.data.DataLoader(testset, batch_size=10000, shuffle=False, num_workers=0)
# 构造迭代器
test_data_iter = iter(testloader)
test_image, test_label = test_data_iter.next()

classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')


# # 展示图片函数
# def imshow(img):
#     img = img / 2 + 0.5  # 相等于unnormalize
#     npimg = img.numpy()
#     # 这里是因为之前transforms.ToTensor()转变为tensor数据将其归一化时改变了通道顺序,这里改回了[H, W, C]
#     plt.imshow(np.transpose(npimg, (1, 2, 0)))
#     plt.show()
#
#
# # print labels
# print(''.join('%7s' % classes[test_label[j]] for j in range(4)))
# # show image
# imshow(torchvision.utils.make_grid(test_image))


# 实例化模型
net = LeNet()
net = net.cuda()

# 定义损失函数
loss_function = nn.CrossEntropyLoss()
# 定义优化器 第一个参数为需要训练的参数，第二个为学习率
optimizer = optim.Adam(net.parameters(), lr=0.001)

# 训练过程
for epoch in range(10):
    running_loss = 0    # 用来垒加训练过程的损失
    for step, data in enumerate(trainloader, start=0):
        # 同时输入一个批次的的图片和标签
        inputs, labels = data
        # 下面这个方法主要用来将一个batch的导数清0，因为计算是一个批次一个批次处理的
        inputs, labels = inputs.cuda(), labels.cuda()   # 修改为gpu,若不用gpu可以注释掉
        optimizer.zero_grad()
        # forward + backward + optimize
        output = net(inputs)
        loss = loss_function(output, labels)
        loss.backward()
        optimizer.step()    # 参数更新

        # print statistics
        running_loss += loss.item()
        if step % 500 == 499:
            with torch.no_grad():
                test_image = test_image.cuda()
                test_label = test_label.cuda()
                outputs = net(test_image)
                predict_y = torch.max(outputs, dim=1)[1]
                accuracy = (predict_y == test_label).sum().item() / test_label.size(0)
                print('[%d, %5d] train_loss:%3f test_accuracy: %3f' %
                      (epoch + 1, step + 1, running_loss/500, accuracy))
                running_loss = 0
print("finish training")
save_path = './LeNet_gpu.pth'
torch.save(net.state_dict(), save_path)

五、prcdict.py

预测文件代码

import torch
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))])
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
net = LeNet()
net.load_state_dict(torch.load('LeNet.pth'))

im = Image.open('bird_test.jpeg')
im = transform(im)
im = torch.unsqueeze(im, dim=0)

with torch.no_grad():
    output = net(im)
    predict = torch.max(output, dim=1)[1].data.numpy()
print(classes[int(predict)])

预测结果：

找到一张bird的图片，可以看到预测正确