pytorch完整训练一个网络的过程

模型搭建部分

#存放一个神经网络架构
import torch
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear


class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        #将上述层全部集合在一个model中
        self.model=Sequential(
            Conv2d(in_channels=3, out_channels=32, kernel_size=5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10)
        )

    def forward(self, x):
        #直接使用集成好的model
        x=self.model(x)
        return x   #不要忘记返回结果

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

训练和测试部分

#训练
import torchvision
#引入在外面搭建好的神经网络架构
from torch.utils.tensorboard import SummaryWriter

from model import *

#实际上第一个参数就是存放数据的文件目录
from torch.utils.data import DataLoader

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)


#查看数据集长度
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)

#引入搭建好的神经网络模块,创建网络模型
mynet=MyNet()


#损失函数——交叉熵损失函数
loss_fn=nn.CrossEntropyLoss()

#优化器
learning_rate=0.01
#需要传入网络的参数到优化器中，因为优化的对象就是网络参数
optimizer=torch.optim.SGD(mynet.parameters(),lr=learning_rate)


#训练网络相关的一些参数
#记录训练的次数
total_train_step=0
#记录测试的次数
total_test_step=0
#训练的轮数
epoch=20

#实际上采用了mini-batch的训练过程
writer=SummaryWriter("logs")
for i in range(epoch):
    print("------第{}轮训练开始-------".format(i))

    #训练步骤开始
    mynet.train()       # 与net.eval()一起管理dropout方法，train()表示启用 BatchNormalization 和 Dropout。 使用model.train() 让model变成训练模式，此时 dropout和batch normalization的操作在训练起到防止网络过拟合的问题。
    for data in train_dataloader:
        imgs,targets=data
        outputs=mynet(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) #使用tensorboard进行记录


    #测试步骤开始
    mynet.eval() #eval()表示不启用 BatchNormalization 和 Dropout。此处一定要使用net.eval()把dropout关掉，因为这里我们的目的是测试训练好的网络，而不是在训练网络，没有必要再dropout和再计算BN的方差和均值(BN使用训练的历史值)
    #经过一个epoch的训练后，测试集开始（也可当作验证集）
    total_test_loss=0       #一个epoch的测试集损失
    total_accurary=0
    for data in test_dataloader:
        imgs,targets=data
        outputs=mynet(imgs)
        loss=loss_fn(outputs,targets)
        total_test_loss=total_test_loss+loss
        #进行准确率的计算 看结果和目标有多少是匹配的
        accurary=(outputs.argmax(1)==targets).sum()  #1是横向
        total_accurary=total_accurary+accurary
    total_test_step=total_test_step+1
    print("测试集上的loss:{}".format(total_test_loss))
    print("测试集上的accurary:{}".format(total_accurary/test_data_size))
    writer.add_scalar("test_loss", total_test_loss, total_test_step)  # 使用tensorboard进行记录
    writer.add_scalar("test_accurary", total_accurary/test_data_size, total_test_step)  # 使用tensorboard进行记录


    #每次训练一个epoch便保存模型
    torch.save(mynet,"mynet_{}.pth".format(i))

writer.close()

效果