【Pytorch】简单的模型训练套路

土堆课的练习

import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

from model import *

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

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)

#创建网络模型
jh=Jh() #创建网络模型之前要先搭建神经网络，搭建好的神经网络放在另一个py文件里

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

#优化器
learning_rate=0.01
optimizer=torch.optim.SGD(jh.parameters(),lr=learning_rate)

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

#添加tensorboard
writer=SummaryWriter("logs")

for i in range(epoch):
    print("第{}轮训练开始".format(i+1))

    #训练步骤开始
    for data in train_dataloader:
        imgs,targets=data
        outputs=jh(imgs)
        loss=loss_fn(outputs,targets) #计算损失值

        #优化器优化模型
        optimizer.zero_grad() #用优化器进行梯度清零
        loss.backward() #反向传播得到每个参数节点的梯度
        optimizer.step() #对参数进行优化
        total_train_step=total_train_step+1
        if total_train_step%100==0:
            print("训练次数：{}, Loss:{}".format(total_train_step, loss.item()))
            writer.add_scalar("train_loss",loss.item(),total_train_step)

    #测试步骤开始
    total_test_loss=0
    with torch.no_grad():
        for data in test_dataloader:
            imgs,targets=data
            outputs=jh(imgs)
            loss=loss_fn(outputs,targets)
            total_test_loss=total_test_loss+loss

    print("整体测试集上的Loss:{}".format(total_test_loss))
    writer.add_scalar("test_loss",total_test_loss,total_test_step)
    total_test_step=total_test_step+1

    torch.save(jh,"jh_{}.pth".format(i)) #保存每一轮训练的模型
    print("模型已保存")
 
writer.close()

 搭建的神经网络放在另一个文件中

import torch
from torch import nn


class Jh(nn.Module):
    def __init__(self):
        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(1024,64),
            nn.Linear(64,10)
        )

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

if __name__ == '__main__':
    jh=Jh()
    input=torch.ones((64,3,32,32))
    output=jh(input)
    print(output.shape)