1. Johngo学长首页
  2. 技术点详细复盘专题
  3. 人工智能

【PyTorch】从头搭建并训练一个神经网络模型（图像分类、CNN）

• 人工智能 • 阅读 59

目录

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

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

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

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

【PyTorch】从头搭建并训练一个神经网络模型（图像分类、CNN）
代码如下，重要地方代码中有注释 
import torch
import cv2
import numpy as np
from torchvision import datasets, transforms
import torchvision

transform = transforms.Compose([transforms.ToTensor(),
                                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)
  1. 定义（搭建）自己的神经网络

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


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

class MyModel(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()

这里解释一下为什么第一个全连接层的输入大小为16x5x5，16是因为最后的卷积层有16个filter（即输出的feature map有16个channel），后面的5×5并不是因为最后卷积层的kernel大小为5×5，而是因为最后的feature map大小为5×5，至于为什么是5×5可以按照下面的图算一下：

【PyTorch】从头搭建并训练一个神经网络模型（图像分类、CNN）
  1. 定义损失函数（Loss Function）和优化器（Optimizer）

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


import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(Net.parameters(), lr=0.001, momentum=0.9)
  1. 训练神经网络

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


epochs = 2
e1 = cv2.getTickCount()
for epoch in range(epochs):
    total_loss = 0.0
    for i, data in enumerate(trainloader, 0):

        inputs, labels = data
        optimizer.zero_grad()

        outputs = Net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        if (i+1) % 1000 == 0:
            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.294
第1个epoch：第 2000次：目前的训练损失loss为：2.087
第1个epoch：第 3000次：目前的训练损失loss为：1.902
第1个epoch：第 4000次：目前的训练损失loss为：1.799
第1个epoch：第 5000次：目前的训练损失loss为：1.709
第1个epoch：第 6000次：目前的训练损失loss为：1.660
第1个epoch：第 7000次：目前的训练损失loss为：1.623
第1个epoch：第 8000次：目前的训练损失loss为：1.586
第1个epoch：第 9000次：目前的训练损失loss为：1.550
第1个epoch：第10000次：目前的训练损失loss为：1.497
第1个epoch：第11000次：目前的训练损失loss为：1.468
第1个epoch：第12000次：目前的训练损失loss为：1.469
第2个epoch：第 1000次：目前的训练损失loss为：1.392
第2个epoch：第 2000次：目前的训练损失loss为：1.382
第2个epoch：第 3000次：目前的训练损失loss为：1.364
第2个epoch：第 4000次：目前的训练损失loss为：1.361
第2个epoch：第 5000次：目前的训练损失loss为：1.344
第2个epoch：第 6000次：目前的训练损失loss为：1.349
第2个epoch：第 7000次：目前的训练损失loss为：1.313
第2个epoch：第 8000次：目前的训练损失loss为：1.327
第2个epoch：第 9000次：目前的训练损失loss为：1.306
第2个epoch：第10000次：目前的训练损失loss为：1.302
第2个epoch：第11000次：目前的训练损失loss为：1.275
第2个epoch：第12000次：目前的训练损失loss为：1.288
用CPU训练总共用时：83.7945317 s

用时83.79秒。

  1. 测试模型结果

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%左右）。并且这个博客的目的在于梳理训练神经网络的大致流程，而非构建一个优秀的模型。

  1. 嫌CPU太慢？换GPU训练并推测试试！

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


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

代码及注释如下：


Net.to(device)
epochs = 2
e1 = cv2.getTickCount()
for epoch in range(epochs):
    total_loss = 0.0
    for i, data in enumerate(trainloader, 0):

        inputs, labels = data

        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()

        outputs = Net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        if (i+1) % 1000 == 0:
            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.231
第1个epoch：第 2000次：目前的训练损失loss为：2.032
第1个epoch：第 3000次：目前的训练损失loss为：1.872
第1个epoch：第 4000次：目前的训练损失loss为：1.745
第1个epoch：第 5000次：目前的训练损失loss为：1.702
第1个epoch：第 6000次：目前的训练损失loss为：1.634
第1个epoch：第 7000次：目前的训练损失loss为：1.603
第1个epoch：第 8000次：目前的训练损失loss为：1.548
第1个epoch：第 9000次：目前的训练损失loss为：1.513
第1个epoch：第10000次：目前的训练损失loss为：1.494
第1个epoch：第11000次：目前的训练损失loss为：1.459
第1个epoch：第12000次：目前的训练损失loss为：1.452
第2个epoch：第 1000次：目前的训练损失loss为：1.396
第2个epoch：第 2000次：目前的训练损失loss为：1.372
第2个epoch：第 3000次：目前的训练损失loss为：1.350
第2个epoch：第 4000次：目前的训练损失loss为：1.361
第2个epoch：第 5000次：目前的训练损失loss为：1.338
第2个epoch：第 6000次：目前的训练损失loss为：1.312
第2个epoch：第 7000次：目前的训练损失loss为：1.322
第2个epoch：第 8000次：目前的训练损失loss为：1.289
第2个epoch：第 9000次：目前的训练损失loss为：1.277
第2个epoch：第10000次：目前的训练损失loss为：1.269
第2个epoch：第11000次：目前的训练损失loss为：1.293
第2个epoch：第12000次：目前的训练损失loss为：1.283
用CPU训练总共用时：74.0288839 s

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

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


correct = 0
total = 0
Net.to(device)

e1 = cv2.getTickCount()
with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        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秒，速度也有提升。

Original: https://blog.csdn.net/qq_44166630/article/details/121718998
Author: SinHao22
Title: 【PyTorch】从头搭建并训练一个神经网络模型（图像分类、CNN）

原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/706543/

转载文章受原作者版权保护。转载请注明原作者出处！

赞 (0)
0

【自取】最近整理的，有需要可以领取学习：



Linux核心资料大放送~

全栈面试题汇总（持续更新&可下载）

一个提高学习100%效率的工具！

【超详细】深度学习面试题目！

LeetCode Python刷题答案下载！

LeetCode Java版刷题答案下载！

LeetCode C++ 版本，抓紧保存！

LeetCode GO语言 刷题答案下载！

大家都在看

亲爱的 Coder【最近整理，可免费获取】👉 最新必读书单  | 👏 面试题下载  | 🌎 免费的AI知识星球