【Pytorch】13.搭建完整的CIFAR10模型

项目源码

已上传至githubCIFAR10Model，如果有帮助可以点个star

简介

在前文【Pytorch】10.CIFAR10模型搭建我们学习了用Module来模拟搭建CIFAR10的训练流程
本节将会加入损失函数，梯度下降，TensorBoard来完整搭建一个训练的模型

基本步骤

搭建神经网络最主要的流程是

• 导入数据集（包括训练集和测试集）
• 创建DataLoader
• 创建自定义的神经网络
• 选择损失函数与梯度下降算法
• 进行n轮训练
• n轮训练完成后通过测试集进行验证
• 引入TensorBoard进行可视化
• 保存每轮训练好的模型
  接下来将逐步拆解这每一个步骤

1.导入数据集

因为我们本文是要训练CIFAR10的模型，所以我们导入CIFAR10的数据集

# 1.创建训练数据集
train_dataset = torchvision.datasets.CIFAR10(root='../dataset', train=True, download=True,
                                             transform=torchvision.transforms.ToTensor())
test_dataset = torchvision.datasets.CIFAR10(root='../dataset', train=False, download=True,
                                            transform=torchvision.transforms.ToTensor())
# 记录数据集大小
train_size = len(train_dataset)
test_size = len(test_dataset)

分别导入训练集与测试集，并且分别记录训练集与测试集的大小
对参数的解释可以看【Pytorch】4.torchvision.datasets的使用这篇文章

2.创建DataLoader

DataLoader主要定义了如何在数据集中取数据的规则，具体讲解可以看【Pytorch】5.DataLoder的使用

# 2.创建dataloader
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=True)

3.创建自定义的神经网络

我们可以在网上搜到CIFAR10的网络模型，通过网络模型来搭建网络，具体可以看【Pytorch】10.CIFAR10模型搭建

import torch
from torch import nn


class CIFAR10Model(nn.Module):
    def __init__(self):
        super(CIFAR10Model, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 5, padding=2)
        self.maxpool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(32, 32, 5, padding=2)
        self.maxpool2 = nn.MaxPool2d(2, 2)
        self.conv3 = nn.Conv2d(32, 64, 5, padding=2)
        self.maxpool3 = nn.MaxPool2d