Pytorch 完整的模型训练套路

Pytorch 完整的模型训练套路

model.py

import torch
from torch import nn

# 搭建神经网络
class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64*4*4, 64),
            nn.Linear(64, 10)
        )


    def forward(self, input):
        input = self.model(input)
        return input


# 主函数
if __name__ == '__main__':
    model = Model()

    # 测试网络正确性, 一般是给定一个确定的输入尺寸, 看输出的尺寸是不是我们想要的
    input = torch.ones((64, 3, 32, 32))     # batch_size=64,代表64张图片  3个通道, CIFAR10 是32*32
    output = model(input)                   # 将输入送入网络中，得到输出output
    print(output.shape)


# 当我们输入64张图片的时候，可以看到返回64行数据，每一行数据上有10个数据， 这10个数据代表每一张图片在我们10个类别中的概率

train.py

# 以 CIFAR10 数据集为例，展示一下完整的模型训练套路，完成对数据集的分类问题

import torch
import torchvision
from torch.utils.tensorboard import SummaryWriter

from model import *         # 注意 model 和 train 在一个文件夹下面

# 准备数据集