【Pytorch】模型训练套路---以CIFAR-10为例

CIFAR-10简介
官网链接：CIFAR-10 and CIFAR-100 datasets

CIFAR-10是一个更接近普适物体的彩色图像数据集。CIFAR-10 是由Hinton 的学生Alex Krizhevsky 和Ilya Sutskever 整理的一个用于识别普适物体的小型数据集。一共包含10 个类别的RGB 彩色图片：飞机（ airplane ）、汽车（ automobile ）、鸟类（ bird ）、猫（ cat ）、鹿（ deer ）、狗（ dog ）、蛙类（ frog ）、马（ horse ）、船（ ship ）和卡车（ truck ）。
每个图片的尺寸为32 × 32 ，每个类别有6000个图像，数据集中一共有50000 张训练图片和10000 张测试图片。

1、与MNIST 数据集中目比， CIFAR-10 真高以下不同点
(1)、CIFAR-10 是3 通道的彩色RGB 图像，而MNIST 是灰度图像。
(2)、CIFAR-10 的图片尺寸为32 × 32 ，而MNIST 的图片尺寸为28 × 28 ，比MNIST 稍大。
(3)、相比于手写字符， CIFAR-10 含有的是现实世界中真实的物体，不仅噪声很大，而且物体的比例、特征都不尽相同，这为识别带来很大困难。直接的线性模型如Softmax 在CIFAR-10 上表现得很差。

模型结构：

【源码】

train.py

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

from model import *

#准备数据集
#训练数据集
train_data=torchvision.datasets.CIFAR10("../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)

#创建网络模型
model=Model()

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

#优化器
learning_rate=0.01
optimizer=torch.optim.SGD(model.parameters(),learning_rate)  #SGD随机梯度下降

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


#Tensorboard
writer=SummaryWriter("./logs_train")

for i in range(epoch):
    print('-------第{}轮训练开始------'.format(i+1))
    #训练开始
    model.train()
    for data in test_dataloader:
        imgs,targets=data
        outputs=model(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)

    #测试
    model.eval()
    total_test_loss=0
    total_accuracy=0
    with torch.no_grad():
        for data in test_dataloader:
            imgs,targets=data
            outputs=model(imgs)
            loss=loss_fn(outputs,targets)
            total_test_loss=total_test_loss+loss.item()
            accuracy=(outputs.argmax(1) == targets).sum()
            total_accuracy=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/test_data_size,total_test_step)
    total_test_step+=1

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

writer.close()

model.py

import torch
from torch import nn

#搭建神经网络(十分类网络)
class Model(nn.Module):
    def __init__(self): #模型初始化
        super(Model,self).__init__()
        self.model=nn.Sequential(
            nn.Conv2d(3,32,kernel_size=5,stride=1,padding=2), #卷积
            nn.MaxPool2d(kernel_size=2), #最大池化
            nn.Conv2d(32,32,kernel_size=5,stride=1,padding=2), #卷积
            nn.MaxPool2d(2), #最大池化
            nn.Conv2d(32,64,kernel_size=5,stride=1,padding=2), #卷积
            nn.MaxPool2d(2), #最大池化
            nn.Flatten(), #摊平 (64*4*4)
            nn.Linear(64*4*4,64), #线性层
            nn.Linear(64,10)
        )

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

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

最后我们将训练的图像在Tensorboard上表现出来

【Pytorch】模型训练套路---以CIFAR-10为例

悦读