【Pytorch】4 训练一个图像分类器(Pytorch官方手册CIFAR10案例)

构建图像分类网络(Pytorch官方手册CIFAR10案例)

引言

对于视觉,我们已经创建了一个叫做 totchvision的包,该包含有支持加载类似Imagenet,CIFAR10,MNIST等公共数据集的数据加载模块 torchvision.datasets和支持加载图像数据数据转换模块 torch.utils.data.DataLoader

CIFAR10数据集,它包含十个类别:‘airplane’, ‘automobile’, ‘bird’, ‘cat’, ‘deer’, ‘dog’, ‘frog’, ‘horse’, ‘ship’, ‘truck’。CIFAR-10 中的图像尺寸为3 * 32 * 32,也就是RGB的3层颜色通道,每层通道内的尺寸为32*32。


训练一个图像分类器主要有以下步骤

  • 使用torchvision加载并归一化CIFAR10的训练和测试数据集
  • 定义一个卷积神经网络
  • 定义一个损失函数
  • 在训练样本数据上训练网络
  • 在测试样本数据上测试网络

第一部分:加载数据集

完整代码:

import torch
import torchvision
import torchvision.transforms as transforms

transform = transforms.Compose(
    [
        transforms.ToTensor(),
        # 归一化 平均值0.5 标准差0.5 后续为什么反归一化为 img = img / 2 + 0.5 原因在这里
        # 这里应该是 batch Normalization 为了使训练更容易 使用归一化
        #                        x - mean
        # Batch Normalization = ----------
        #                          std
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ]
)

transet = torchvision.datasets.CIFAR10(root='./data/train', train=True, download=True, transform=transform)
testset = torchvision.datasets.CIFAR10(root='./data/test', train=False, download=True, transform=transform)

trainloader = torch.utils.data.DataLoader(transet, batch_size=4, shuffle=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False)

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

注释一:
transforms模块主要包括对图像的一些操作:

详情可以参考这位大佬的博客:torchvision库及其常用函数

transform = transforms.Compose(
    [
        transforms.ToTensor(),
        # 归一化 平均值0.5 标准差0.5 后续为什么反归一化为 img = img / 2 + 0.5 原因在这里
        # 这里应该是 batch Normalization 为了使训练更容易 使用归一化
        #                        x - mean
        # Batch Normalization = ----------
        #                          std
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ]
)

注释二:加载数据集操作

torchvision.datasets.*用于加载常用数据集,Pytorch目前支持31种数据集:

- Caltech、CelebA、CIFAR、Cityscapes、COCO、EMNIST
- FakeData、Fashion-MNIST、Flickr、HMDB51、ImageNet、iNaturalist
- Kinetics-400、KITTI、KMNIST、LFW、LSUN
- MNIST、Omniglot、PhotoTour、Places365、QMNIST
- SBD、SBU、SEMEION、STL10、SVHN
- UCF101、USPS、VOC、WIDERFace、Base classes for custom datasets

torchvision.datasets.CIFAR10参数:

  • root:数据集下载到的根目录,数据集所在的根目录,这里是cifar-10-batches-py的根目录
  • trainTrue代表加载训练集,False代表加载测试集
  • downloadTrue从互联网上下载数据,并将其放在root目录下,如果数据集已经下载,什么也不干。(这里下载可能速度慢,需要魔法)链接:CIFAR
    提取码:8lwp
  • trasnform:设置一组对图像进行处理的操作,这一组操作由Compose组成,这一组compose 的顺序还很重要。

torchvision.data.Dataloader参数:

  • trainset:将刚刚生成的trainset列表传入torchvision.data.Dataloader的第一个参数。
  • batch size:设定图像批次大小,就是一次处理几张图片
  • shuffleTrue代表每一个epoch过程中会打乱数据顺序,重新随机选择,一般都会设置为True
  • num worker:导入数据时用几个线程操作,windows一般设置为0,我用的windows设置为2时程序报错,Linux可以设成2.
transet = torchvision.datasets.CIFAR10(root='./data/train', train=True, download=True, transform=transform)
testset = torchvision.datasets.CIFAR10(root='./data/test', train=False, download=True, transform=transform)

trainloader = torch.utils.data.DataLoader(transet, batch_size=4, shuffle=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False)

第二部分:训练数据

完整代码:

class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()

        self.conv1 = nn.Conv2d(3, 6, 5)
        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 = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, self.num_flat_features(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

device = torch.device('cpu:0')
# 第一步:创建网络
net = Net()
# net.to(device)
# 第二部:创建损失函数 这里用的交叉熵
criterion = nn.CrossEntropyLoss()
# criterion.to(device)
# 第三步:创建权重更新规则 这里用的SGD
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# 第四步:循环迭代数据 在数据迭代器上循环传给网络和优化器输入
since = time.time()
for epoch in range(2):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        # inputs = inputs.to(device)
        # labels = labels.to(device)

        # 核心代码
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 2000 == 1999:
            print('[{0},{1}] loss:{2}'.format(epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
    time_elapsed // 60, time_elapsed % 60))
print('Finished Training')

注释一:创建网络、损失函数、设置权重更新规则

# 第一步:创建网络
net = Net()
# 第二部:创建损失函数 这里用的交叉熵
criterion = nn.CrossEntropyLoss()
# 第三步:创建权重更新规则 这里用的SGD
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

注释二:循环迭代数据

# 迭代2次epoch

以上都是必要步骤,核心代码。

注释三:可能会经常用到的代码

  1. 更新权重核心代码
# 核心代码
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
  1. 计算训练时间核心代码
import time

since = time.time()

for epoch in range(2):
    ···
    
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
    time_elapsed // 60, time_elapsed % 60))
  1. 显示图像
def imshow(img):
    # 反归一化
    img = img / 2 + 0.5
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

    ···
    
# 方法一:官方实例
dataiter = iter(trainloader)
images, labels = dataiter.next()
imshow(torchvision.utils.make_grid(images))
print(' '.join('%5s' % classes[labels[j]] for j in range(4)))

# 方法二:网上看到的
for image, label in trainloader:
    imshow(image[0])

注释四:关于GPU加速

关于如何使用GPU加速,主要对Net``````criterion``````image``````labels进行to(device)操作即可。

也可以使用.cuda(),例如net.cuda(),但看网上博客说不推荐使用,尽量使用to(device)

我用的GPU是RTX3060,目前使用GPU加速与不使用GPU加速我测试的训练时间基本一样,稳定在52s,没有体现出来GPU加速的效果。我猜测可能还是因为网络比较简单或者数据集数量比较小。或者还是我的代码有错误,如果有错误,希望大佬们可以评论指出。

另外一个是我用RTX3060遇到了显卡不兼容的问题,很有可能是Pytorch版本太低了,我用的CUDA是11.3,但在之前使用的Pytorch是1.9.0版本,升级到Pytorch 1.10.0就顺利解决问题了。
参考博客:显卡不兼容问题

我直接搬运过来了,如果你也是RTX3060,直接在Pytorch所在环境下运行下面代码即可:

# 使用conda卸载Pytorch 
conda uninstall pytorch
conda uninstall libtorch

# 使用pip卸载Pytorch
pip uninstall torch

# 安装Pytorch 1.10.0 + gpu版本 需要CUDA 11.3版本
pip3 --default-timeout=60000 install torch==1.10.0+cu113 torchvision==0.11.1+cu113 torchaudio===0.10.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html -i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

net = Net()
net.to(device)

criterion = nn.CrossEntropyLoss()
criterion.to(device)

··

for epoch in range(2):
    ···
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        inputs = inputs.to(device)
        labels = labels.to(device)
    ···

第三部分:程序完整代码

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

import torchvision
import torchvision.transforms as transforms

import matplotlib.pyplot as plt
import numpy as np

import torch.optim as optim

import time


class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()

        self.conv1 = nn.Conv2d(3, 6, 5)
        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 = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, self.num_flat_features(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]
        num_features = 1
        for s in size:
            num_features *= s
        return num_features


def imshow(img):
    # 反归一化
    img = img / 2 + 0.5
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


transform = transforms.Compose(
    [
        transforms.ToTensor(),
        # 归一化 平均值0.5 标准差0.5 后续为什么反归一化为 img = img / 2 + 0.5 原因在这里
        # 这里应该是 batch Normalization 为了使训练更容易 使用归一化
        #                        x - mean
        # Batch Normalization = ----------
        #                          std
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ]
)

transet = torchvision.datasets.CIFAR10(root='./data/train', train=True, download=True, transform=transform)
testset = torchvision.datasets.CIFAR10(root='./data/test', train=False, download=True, transform=transform)

trainloader = torch.utils.data.DataLoader(transet, batch_size=4, shuffle=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False)

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

# 展示一些图片
# # get some random training images
# dataiter = iter(trainloader)
# images, labels = dataiter.next()
#
# # show images
# imshow(torchvision.utils.make_grid(images))
# # print labels
# print(' '.join('%5s' % classes[labels[j]] for j in range(4)))


# for image, label in trainloader:
#     imshow(image[0])

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

# 第一步:创建网络
net = Net()
# net.to(device)
# 第二部:创建损失函数 这里用的交叉熵
criterion = nn.CrossEntropyLoss()
# criterion.to(device)
# 第三步:创建权重更新规则 这里用的SGD
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# 第四步:循环迭代数据 在数据迭代器上循环传给网络和优化器输入
since = time.time()
for epoch in range(2):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        # inputs = inputs.to(device)
        # labels = labels.to(device)

        # 核心代码
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 2000 == 1999:
            print('[{0},{1}] loss:{2}'.format(epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
    time_elapsed // 60, time_elapsed % 60))
print('Finished Training')

你可能感兴趣的:(Pytorch学习笔记,pytorch,深度学习,计算机视觉)