在pytroch中使用CIFAR10完成完整的模型训练套路

模型训练套路：

• 1.准备数据集
• 2.加载数据集
• 3.搭建神经网络
• 4创建损失函数
• 5.优化器
• 6.设置训练网络的一些参数
• 7.添加tensorboard（方便观察）
• 8.开始训练
• .测试
• 9.保存神经网络

准备数据

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

加载数据

#加载数据集
loader_train=DataLoader(dataset_train,batch_size=16,drop_last=True,shuffle=False)
loader_test=DataLoader(dataset_test,batch_size=16,drop_last=True,shuffle=False)

搭建神经网络

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.model1=torch.nn.Sequential(
            nn.Conv2d(3, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(1024, 64),
            nn.Linear(64, 10)
        )

    def forward(self,x):
        return self.model1(x)

# #验证神经网络正确性
#
if __name__=="__main__":
    net = Net()
    input=torch.ones((64,3,32,32))
    output=net(input)
    print(output.shape)

创建损失函数

#创建损失函数
loss_cro=torch.nn.CrossEntropyLoss()

优化器

#创建优化器
optim=torch.optim.SGD(net.parameters(),lr=1e-3)

设置训练网络的一些参数

#设置训练网络的一些参数
total_train_stp=0#总训练次数
total_test_stp=0#总测试次数
epoch=20#训练轮次

添加tensorboard（方便观察）

#tensorboard
writer=SummaryWriter("./end")

开始训练

for i in range(epoch):
    loss=0
    print("-----------第{}轮训练开始--------------".format(i+1))
    for data in loader_train:
        imgs,targets=data
        output=net(imgs)
        loss=loss_cro(output,targets)

        #优化器优化
        optim.zero_grad()
        loss.backward()
        optim.step()

        #记录训练次数
        total_train_stp+=1
        if total_train_stp%100==0:
            print("训练次数：{}，loss:{}".format(total_train_stp,loss.item()))
            writer.add_scalar("train_loss",loss.item(),total_train_stp)

测试

#测试
with torch.no_grad():
    total_current_test=0
    for data in loader_test:
        imgs,targets=data
        output=net(imgs)
        loss=loss_cro(output,targets)
        total_test_stp+=1

        writer.add_scalar("test_loss",loss,total_test_stp)
        accuracy=(output.argmax(1)==targets).sum()
        total_current_test+=accuracy
        print("整体测试的正确率为{}".format(total_current_test/(total_test_stp)*10))
        writer.add_scalar("test_accuracy",total_current_test/(total_test_stp)*10,total_test_stp)

保存模型

torch.save(net,"./model3.pth")

总代码

总代码分在两个py文件中，分别是

• 负责训练模型的train.py
• 神经网络所在文件network.py

network.py

import torch
from torch import nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.model1=torch.nn.Sequential(
            nn.Conv2d(3, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(1024, 64),
            nn.Linear(64, 10)
        )

    def forward(self,x):
        return self.model1(x)

# #验证神经网络正确性
#
if __name__=="__main__":
    net = Net()
    input=torch.ones((64,3,32,32))
    output=net(input)
    print(output.shape)

train.py

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

from network import Net


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


#加载数据集
loader_train=DataLoader(dataset_train,batch_size=16,drop_last=True,shuffle=False)
loader_test=DataLoader(dataset_test,batch_size=16,drop_last=True,shuffle=False)

#搭建神经网络
net=Net()


#创建损失函数
loss_cro=torch.nn.CrossEntropyLoss()

#创建优化器
optim=torch.optim.SGD(net.parameters(),lr=1e-3)

#设置训练网络的一些参数
total_train_stp=0#总训练次数
total_test_stp=0#总测试次数
epoch=20#训练轮次

#tensorboard
writer=SummaryWriter("./end")

#开始训练
for i in range(epoch):
    loss=0
    print("-----------第{}轮训练开始--------------".format(i+1))
    for data in loader_train:
        imgs,targets=data
        output=net(imgs)
        loss=loss_cro(output,targets)

        #优化器优化
        optim.zero_grad()
        loss.backward()
        optim.step()

        #记录训练次数
        total_train_stp+=1
        if total_train_stp%100==0:
            print("训练次数：{}，loss:{}".format(total_train_stp,loss.item()))
            writer.add_scalar("train_loss",loss.item(),total_train_stp)

#测试
with torch.no_grad():
    total_current_test=0
    for data in loader_test:
        imgs,targets=data
        output=net(imgs)
        loss=loss_cro(output,targets)
        total_test_stp+=1

        writer.add_scalar("test_loss",loss,total_test_stp)
        accuracy=(output.argmax(1)==targets).sum()
        total_current_test+=accuracy
        print("整体测试的正确率为{}".format(total_current_test/(total_test_stp)*10))
        writer.add_scalar("test_accuracy",total_current_test/(total_test_stp)*10,total_test_stp)
torch.save(net,"./model3.pth")
print("模型已经保存")

writer.close()

最终运行截图

可以看出再训练了20轮后可以达到93%的正确率