PyTorch从头搭建并训练一个神经网络模型(图像分类CNN)

Posted Rainbowman 0

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了PyTorch从头搭建并训练一个神经网络模型(图像分类CNN)相关的知识,希望对你有一定的参考价值。

0. 前言

之前用过一些很厉害的模型,图像分类领域的VGG16,目标检测领域的YoloV5,实例分割领域的Yolact等。但只是会配置好环境之后训练,最多稍微修改下源码的接口满足自己的需求。还从来没有用PyTorch从头搭建并训练一个模型出来。

正好最近在较为系统地学PyTorch,就总结一下如何从头搭建并训练一个神经网络模型。

1. 使用torchvision加载数据集并做预处理

我们使用的数据集是CIFAR10,该数据集有10个类别,图像尺寸为3 x 32 x 32,如下所示:

代码如下,重要地方代码中有注释

import torch
import cv2
import numpy as np
from torchvision import datasets, transforms
import torchvision

# 1. 使用torchvision加载数据集并做预处理
transform = transforms.Compose([transforms.ToTensor(),  # 把图像转换为tensor
                                transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5)),  # 归一化处理
                              ])
# 加载训练集和测试集
trainset = torchvision.datasets.CIFAR10(root='E:\\\\Machine Learning\\\\PyTorch\\\\CIFAR10', 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='E:\\\\Machine Learning\\\\PyTorch\\\\CIFAR10', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2)

2. 定义(搭建)自己的神经网络

代码及注释如下,整个结构很简单,就是两个卷积层,两个最大池化层,最后连接三个全连接层。

# 2. 定义卷积神经网络
import torch.nn as nn
import torch.nn.functional as F

class MyModel(nn.Module):  #  继承nn.Module
    # 定义网络结构
    def __init__(self):
        super(MyModel, 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

Net = MyModel()

3. 定义损失函数(Loss Function)和优化器(Optimizer)

神经网络的反向传播需要损失函数,因为是多分类问题,所以我们用交叉熵损失函数:

# 3. 定义损失函数和优化器
import torch.optim as optim

criterion = nn.CrossEntropyLoss()  # 多分类问题,用交叉熵损失函数
optimizer = optim.SGD(Net.parameters(), lr=0.001, momentum=0.9)  # 用SGD优化器

4. 训练神经网络

接着开始训练我们的模型,共训练2个epoch,batch_size=4,先用CPU训练,看看时间如何。

# 4. 训练神经网络
epochs = 2  # 训练两个epoch,batch_size = 4 (batch_size的大小定义在第一步torch.utils.data.DataLoader中)
e1 = cv2.getTickCount()  # 记录训练时间
for epoch in range(epochs):
    total_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # 得到inputs
        inputs, labels = data
        optimizer.zero_grad()
        # forward + backward + optimize
        # 前向传播+反向传播+更新参数
        outputs = Net(inputs)  # 前向传播,得到outputs
        loss = criterion(outputs, labels)  # 得到损失函数
        loss.backward()  # 后向传播
        optimizer.step()  # 更新参数

        # 输出训练过程
        total_loss += loss.item()
        if (i+1) % 1000 == 0:  # 每1000次(就是4000张图像)输出一次loss
            print('第个epoch:第:5d次:目前的训练损失loss为::.3f'.format(epoch+1, i+1, total_loss/1000))
            total_loss = 0.0

e2 = cv2.getTickCount()
print('用CPU训练总共用时: s'.format((e2-e1)/cv2.getTickFrequency()))
# 输出结果:1个epoch:第 1000次:目前的训练损失loss为:2.2941个epoch:第 2000次:目前的训练损失loss为:2.0871个epoch:第 3000次:目前的训练损失loss为:1.9021个epoch:第 4000次:目前的训练损失loss为:1.7991个epoch:第 5000次:目前的训练损失loss为:1.7091个epoch:第 6000次:目前的训练损失loss为:1.6601个epoch:第 7000次:目前的训练损失loss为:1.6231个epoch:第 8000次:目前的训练损失loss为:1.5861个epoch:第 9000次:目前的训练损失loss为:1.5501个epoch:第10000次:目前的训练损失loss为:1.4971个epoch:第11000次:目前的训练损失loss为:1.4681个epoch:第12000次:目前的训练损失loss为:1.4692个epoch:第 1000次:目前的训练损失loss为:1.3922个epoch:第 2000次:目前的训练损失loss为:1.3822个epoch:第 3000次:目前的训练损失loss为:1.3642个epoch:第 4000次:目前的训练损失loss为:1.3612个epoch:第 5000次:目前的训练损失loss为:1.3442个epoch:第 6000次:目前的训练损失loss为:1.3492个epoch:第 7000次:目前的训练损失loss为:1.3132个epoch:第 8000次:目前的训练损失loss为:1.3272个epoch:第 9000次:目前的训练损失loss为:1.3062个epoch:第10000次:目前的训练损失loss为:1.3022个epoch:第11000次:目前的训练损失loss为:1.2752个epoch:第12000次:目前的训练损失loss为:1.288
用CPU训练总共用时:83.7945317 s

用时83.79秒。

5. 测试模型结果

# 5. 测试模型准确率如何

correct = 0
total = 0

e1 = cv2.getTickCount()  # 记录测试时间
with torch.no_grad():  # 不跟踪梯度
    for data in testloader:
        images, labels = data
        outputs = Net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        
e2 = cv2.getTickCount()
print('用CPU测试总共用时: s'.format((e2-e1)/cv2.getTickFrequency()))
print('在测试集上的准确率为::.3f%'.format(correct*100/total))

我们来看看结果如何

# 结果
用CPU测试总共用时:5.3146039 s
在测试集上的准确率为:55.460%

准确率大概55%,虽然不高,但可以看出我们的模型确实学到了东西(因为10分类问题随机预测的话准确率在10%左右)。并且这个博客的目的在于梳理训练神经网络的大致流程,而非构建一个优秀的模型。

6. 嫌CPU太慢?换GPU训练并推测试试!

首先看看是否安装了对应版本的cuda和cudnn,具体安装步骤就不说了,csdn上很多优秀教程。

# 检测电脑上是否安装了对应版本的cuda
device = torch.device('cuda:0' if torch.cuda.is_available else 'cpu')
print('设备名称: ', device)
print('查看cuda版本: ', torch.version.cuda)
# 结果
设备名称:  cuda:0
查看cuda版本:  10.1

PyTorch使用GPU训练非常方便,相较于第4步用CPU训练,只需增加两行代码:
(1)把神经网络模型加载到cuda
(2)把数据加载到cuda

代码及注释如下:

# 6. 用GPU训练神经网络
######  第1处不一样   ########
Net.to(device)  # 把神经网络模型加载到cuda
epochs = 2  # 训练两个epoch,batch_size = 4 (batch_size的大小定义在第一步torch.utils.data.DataLoader中)
e1 = cv2.getTickCount()  # 记录训练时间
for epoch in range(epochs):
    total_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # 得到inputs
        inputs, labels = data
        ######  第2处不一样   ########
        inputs, labels = inputs.to(device), labels.to(device)  # 把数据加载到cuda
        optimizer.zero_grad()
        # forward + backward + optimize
        # 前向传播+反向传播+更新参数
        outputs = Net(inputs)  # 前向传播,得到outputs
        loss = criterion(outputs, labels)  # 得到损失函数
        loss.backward()  # 后向传播
        optimizer.step()  # 更新参数

        # 输出训练过程
        total_loss += loss.item()
        if (i+1) % 1000 == 0:  # 每1000次(就是4000张图像)输出一次loss
            print('第个epoch:第:5d次:目前的训练损失loss为::.3f'.format(epoch+1, i+1, total_loss/1000))
            total_loss = 0.0

e2 = cv2.getTickCount()
print('用CPU训练总共用时: s'.format((e2-e1)/cv2.getTickFrequency()))

看看结果如何:

# 结果1个epoch:第 1000次:目前的训练损失loss为:2.2311个epoch:第 2000次:目前的训练损失loss为:2.0321个epoch:第 3000次:目前的训练损失loss为:1.8721个epoch:第 4000次:目前的训练损失loss为:1.7451个epoch:第 5000次:目前的训练损失loss为:1.7021个epoch:第 6000次:目前的训练损失loss为:1.6341个epoch:第 7000次:目前的训练损失loss为:1.6031个epoch:第 8000次:目前的训练损失loss为:1.5481个epoch:第 9000次:目前的训练损失loss为:1.5131个epoch:第10000次:目前的训练损失loss为:1.4941个epoch:第11000次:目前的训练损失loss为:1.4591个epoch:第12000次:目前的训练损失loss为:1.4522个epoch:第 1000次:目前的训练损失loss为:1.3962个epoch:第 2000次:目前的训练损失loss为:1.3722个epoch:第 3000次:目前的训练损失loss为:1.3502个epoch:第 4000次:目前的训练损失loss为:1.3612个epoch:第 5000次:目前的训练损失loss为:1.3382个epoch:第 6000次:目前的训练损失loss为:1.3122个epoch:第 7000次:目前的训练损失loss为:1.3222个epoch:第 8000次:目前的训练损失loss为:1.2892个epoch:第 9000次:目前的训练损失loss为:1.2772个epoch:第10000次:目前的训练损失loss为:1.2692个epoch:第11000次:目前的训练损失loss为:1.2932个epoch:第12000次:目前的训练损失loss为:1.283
用CPU训练总共用时:74.0288839 s

相较于CPU的83秒,GPU用了74秒,快了一些,但提升不够明显。这是因为我们的网络很小,参数也很少。另外我的笔记本的GPU也挺老的,1050核显。

我们在试试用GPU推断会加速多少:

# 6. 使用GPU推断

correct = 0
total = 0
Net.to(device)  # 把神经网络模型加载到cuda

e1 = cv2.getTickCount()  # 记录测试时间
with torch.no_grad():  # 不跟踪梯度
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device) # 把数据加载到cuda
        outputs = Net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        
e2 = cv2.getTickCount()
print('用CPU测试总共用时: s'.format((e2-e1)/cv2.getTickFrequency()))
print('在测试集上的准确率为::.3f%'.format(correct*100/total))

看看结果如何

# 结果
用CPU测试总共用时:4.3125164 s
在测试集上的准确率为:54.090%

相较于CPU推断需要5.3秒,使用GPU推断需要4.3秒,速度也有提升。

以上是关于PyTorch从头搭建并训练一个神经网络模型(图像分类CNN)的主要内容,如果未能解决你的问题,请参考以下文章

图像分类---利用pytorch搭建AlexNet网络模型训练自己的数据集(猫狗分类)

奉献pytorch 搭建 CNN 卷积神经网络训练图像识别的模型,配合numpy 和matplotlib 一起使用调用 cuda GPU进行加速训练

「深度学习一遍过」必修5:从头训练自己的数据无从下手?来看看这10个pytorch自带的分类模型叭

从头开始训练 Resnet 深度神经网络

小白学习PyTorch教程六基于CIFAR-10 数据集,使用PyTorch 从头开始​​构建图像分类模型

Pytorch定义并训练自己的数字数据集