Pytorch学习笔记--完整模型的训练套路-优快云博客

每块代码作用在注释中已经给出：

#准备数据集
import torch
import torch.nn as nn  # 替换掉原来的 import nn
import torchvision
from keras.src.metrics.accuracy_metrics import accuracy
from torch.nn import Conv2d, MaxPool2d, Sequential
from torch.utils.data import DataLoader
from model import *
from torch.utils.tensorboard import  SummaryWriter
import numpy as np
from PIL import Image  #注意这里是PIL型。我们要把图片转换成numpy型


#准备数据集
train_data=torchvision.datasets.CIFAR10(root="../data",train=True,transform=torchvision.transforms.ToTensor(),download=True)
test_data=torchvision.datasets.CIFAR10(root="../data",train=False,transform=torchvision.transforms.ToTensor(),download=True)

#length 长度
train_data_size=len(train_data)
test_data_size=len(test_data)

print("训练数据集的长度为:{}".format(train_data_size))
print("测试数据集的长度为:{}".format(test_data_size))

#利用Dataloader加载数据集
train_dataloader=DataLoader(train_data,batch_size=64)
test_dataloader=DataLoader(test_data,batch_size=64)

#创建模型
aa=Aaax()

#损失函数
loss_fn=nn.CrossEntropyLoss()

#优化器
learning_rate=1e-2  #（0.01）
optimizer=torch.optim.SGD(aa.parameters(),lr=learning_rate)

#设置训练网络的一些参数
total_train_step=0  #记录训练的次数
total_test_step=0 #记录测试的次数
epoch=10  #训练的轮数

#添加tensorboard
writer=SummaryWriter("../logs_train")


for i in range(epoch):
    print("------第{}轮训练开始------".format(i+1))
    #训练步骤开始
    for data in train_dataloader:
        imgs,targets=data
        output=aa(imgs)
        loss=loss_fn(output,targets)

        #优化器优化模型
        optimizer.zero_grad()  #梯度清零
        loss.backward()  #反向传播，得到每个参数节点的梯度
        optimizer.step()  #对参数优化
        total_train_step+=1
        if total_train_step%100==0: # 逢百打印
            print("训练次数：{}，loss：{}".format(total_train_step,loss))
            writer.add_scalar("train_loss",loss.item(),total_train_step)
    #测试步骤开始 :在测试集上再跑一遍，通过看测试集上的损失来评估模型的性能
    total_test_loss=0
    total_accuracy=0
    with torch.no_grad():       #测试时要注意不要有梯度！！！
        for data in test_dataloader:
            imgs,targets=data
            output=aa(imgs)
            loss=loss_fn(output,targets)
            total_test_loss+=loss.item()
            #计算准确率
            accuracy=(output.argmax(1)==targets).sum()  #计算正确的个数之和
            total_accuracy+=accuracy
    print("整体测试集上的loss：{}".format(total_test_loss))
    print("整体测试集上的准确率：{}".format(total_accuracy/test_data_size))
    writer.add_scalar("test_loss",total_test_loss,total_test_step)
    writer.add_scalar("test_accuracy",total_accuracy,total_test_step)
    total_test_step+=1

    torch.save(aa,"aa_{}.pth".format(i))
    print("模型已保存")

writer.close()

模型：

import torch
from torch import nn

# 搭建神经网络 CIFAR-10 有十个类别，所以要写一个十分类的神经网络
class Aaax(nn.Module):  # 将 nn.module 改为 nn.Module
    def __init__(self):
        super(Aaax, self).__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, x):
        x = self.model(x)
        return x

# if __name__ == '__main__':
#     aa = Aaax()
#     input = torch.ones((64, 3, 32, 32))  # batch size, channel, 32*32 dimensions
#     output = aa(input)
    # print(output.shape)