前言

上篇博文整理了如何用Pytorch搭建一个基本网络模型，本篇进行一个图像分类任务实操。
相关代码主要参考自官网教程：
https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html#sphx-glr-beginner-blitz-cifar10-tutorial-py

数据集简介

CIFAR10数据集总共包含10个类别，每张图片为3通道的RGB图片，大小为32x32像素。

数据集下载与预处理

使用torchvision.datasets可以下载经典数据集，设置下载路径root和download=True，电脑会自动检测root路径下有无数据集，若不存在数据集，则自动进行下载。

数据预处理使用到了torchvision.transforms这个工具包。

• transforms.Compose
  能够将多个预处理步骤进行合并
• transforms.ToTensor()
  将数据归一化到(0,1)，原始图像数据范围是(0,255)，归一化有利于加快运算
• transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
  对数据进行标准化，公式是(x-mean)/std，前一个是三通道的均值，后一个是三通道的方差，处理之后数值范围为(-1,1)

经过处理后，数据服从正态分布。

import torch
import torchvision
import torchvision.transforms as transforms

transform = transforms.Compose(
    [transforms.ToTensor() # 归一化到(0,1)
     ,transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])  # 再 (x-mean)/std，归一化到(-1,1)
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                          shuffle=True, num_workers=2)
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                         shuffle=False, num_workers=2)

定义卷积神经网络

卷积层 nn.Conv2d，最大池化层nn.MaxPool2d，数值不做具体分析。
注意里面有个view的小技巧x.view(-1, 16 * 5 * 5),-1代表让程序自行进行计算填充。

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


class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

定义损失函数

使用SGD优化器，交叉熵损失

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(mynet.parameters(), lr=0.001, momentum=0.9)

训练

使用CPU进行训练

# 在 Jupyter 中的进度条
from tqdm.notebook import tqdm
import torch.optim as optim

n_epochs = 2
n_batches = 100
mynet = MyNet()
model_path = "cifar_net.pth"
mynet.load_state_dict(torch.load(model_path))
pbar = tqdm(total=n_epochs)
# 使用训练数据训练 n_epochs 轮
for epoch in range(n_epochs):
    # 当前 n_batches 个小批次训练数据上的平均损失
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # 获取输入数据，这里 data 是形如 [inputs, labels] 的列表
        inputs, labels = data
        # 对梯度进行清零操作，防止错误的梯度累加
        optimizer.zero_grad()
        # 前向传播
        outputs = mynet(inputs)
        # 计算损失
        loss = criterion(outputs, labels)
        # 反向传播并更新模型参数
        loss.backward()
        optimizer.step()
        # 统计损失函数的滑动平均
        running_loss += loss.item()
        if (i+1) % n_batches == 0:
            pbar.set_postfix(batch=i+1, loss=running_loss/n_batches)
            running_loss = 0.0
    pbar.update()  # 更新进度条
print('训练完成')
model_path = "cifar_net_100.pth"
torch.save(mynet.state_dict(), model_path)

调用GPU进行训练

调用GPU训练很简单，首先写这句device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")来判断电脑上的GPU是否好用，如果可以用就调用第0块GPU。后面再将数据和模型.to(device)放置到GPU上就可以了。
在老版本中，也看到.cuda()这种写法，效果一样。
GPU版本训练代码如下，可和上面这段进行对比。

# 在 Jupyter 中的进度条
from tqdm.notebook import tqdm
import torch.optim as optim

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

n_epochs = 2
n_batches = 100
mynet = MyNet()
mynet = mynet.to(device)
model_path = "cifar_net.pth"
mynet.load_state_dict(torch.load(model_path))
pbar = tqdm(total=n_epochs)
# 使用训练数据训练 n_epochs 轮
for epoch in range(n_epochs):
    # 当前 n_batches 个小批次训练数据上的平均损失
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # 获取输入数据，这里 data 是形如 [inputs, labels] 的列表
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)
        # 对梯度进行清零操作，防止错误的梯度累加
        optimizer.zero_grad()
        # 前向传播
        outputs = mynet(inputs)
        # 计算损失
        loss = criterion(outputs, labels)
        # 反向传播并更新模型参数
        loss.backward()
        optimizer.step()
        # 统计损失函数的滑动平均
        running_loss += loss.item()
        if (i+1) % n_batches == 0:
            pbar.set_postfix(batch=i+1, loss=running_loss/n_batches)
            running_loss = 0.0
    pbar.update()  # 更新进度条
print('训练完成')
model_path = "cifar_net_100.pth"
torch.save(mynet.state_dict(), model_path)

我的显卡版本是RTX2060，在Jupyter中，可以使用%%time语句来统计cell的执行时间。实测下来，GPU和CPU训练速度并没有明显差异，个人猜测可能是由于该数据量不大的原因。

测试

整体测试

加载模型，在测试集上评估模型性能

mynet = MyNet()
model_path = "cifar_net_100.pth"
mynet.load_state_dict(torch.load(model_path))
correct = 0
total = 0
# 评估时无需计算梯度
with torch.no_grad():
    for data in testloader:
        # 获取数据集中的图像及对应标签真值
        images, labels = data
        # 计算网络的预测的图像概率向量
        outputs = mynet(images)
        # 根据网络输出概率获得预测标签 torch.max
        _, predicted = torch.max(outputs.data, 1)
        # 累计样本总数
        total += labels.size(0)
        # 累加正确预测的样本数
        correct += (predicted == labels).sum().item()
print('网络在测试集上的准确率是: %3.1f %%' % (100*correct/total))

各类样本测试

统计各类别测试样本准确率

class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = mynet(images)
        _, predicted = torch.max(outputs, 1)
        c = (predicted == labels).squeeze()
        for i in range(4):
            label = labels[i]
            class_correct[label] += c[i].item()
            class_total[label] += 1
for i in range(10):
    print('Accuracy of %5s : %2d %%' % (
        classes[i], 100 * class_correct[i] / class_total[i]))

运行结果：

文章来源: zstar.blog.csdn.net，作者：zstar-_，版权归原作者所有，如需转载，请联系作者。

原文链接：zstar.blog.csdn.net/article/details/125752073