python

关注公众号 jb51net

关闭
首页 > 脚本专栏 > python > pytorch中的模型训练

pytorch中的模型训练(以CIFAR10数据集为例)

作者:MarkAssassin

这篇文章主要介绍了pytorch中的模型训练(以CIFAR10数据集为例),具有很好的参考价值,希望对大家有所帮助。如有错误或未考虑完全的地方,望不吝赐教

在pytorch模型训练时,基本的训练步骤可以大致地归纳为:

准备数据集--->搭建神经网络--->创建网络模型--->创建损失函数--->设置优化器--->训练步骤开始--->测试步骤开始

本文以pytorch官网中torchvision中的CIFAR10数据集为例进行讲解。

需要用到的库为(这里说一个小技巧,比如可以在没有import对应库的情况下先输入"torch",,之后将光标移到torch处单击,这时左边就会出现一个红色的小灯泡,点开它就可以import对应的库了):

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

准备数据集

数据集分为“训练数据集”+“测试数据集”。

CIFAR数据集是由50000训练集和10000测试集组成。

这里可以调用torchvision.datasets对CIFAR10数据集进行获取:

#训练数据集
train_data=torchvision.datasets.CIFAR10(root='./dataset',train=True,transform=torchvision.transforms.ToTensor(),download=True)
#测试数据集
test_data=torchvision.datasets.CIFAR10(root='./dataset',train=False,transform=torchvision.transforms.ToTensor(),download=True)
#利用dataloader来加载数据集
train_data_loader=DataLoader(train_data,batch_size=64)
test_data_loader=DataLoader(test_data,batch_size=64)

root是数据集保存的路径,这里笔者使用的是相对路径;对于训练集train=True,而测试集train=False;transform是将数据集的类型转换为tensor类型;download一般设置为True。

之后利用DataLoader对数据集进行加载即可,其中batch_size表示单次传递给程序用以训练的数据(样本)个数。

(这里可以再获取下测试集数据的长度,这样后面在测试步骤时就可以通过计算得到整体测试集上的准确率

搭建神经网络

pytorch官网提供了一个神经网络的简单实例:

import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)
    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

这个实例展示了神经网络的基本架构。

从百度上我们可以搜索到CIFAR10数据集的网络基本架构:

 从原理图中可以看到,该网络从输入(inputs)到输出(outputs)先后经过了

卷积(Convolution)--->最大池化(Max-pooling)--->卷积--->最大池化--->卷积--->最大池化--->展平(Flatten)--->2次线性层

由此可以开始搭建神经网络,笔者将网络命名为MXC:

class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)
#搭建神经网络
class MXC(nn.Module):
    def __init__(self):
        super(MXC, 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),
            Flatten(),
            nn.Linear(64*4*4,64),
            nn.Linear(64,10)
        )
    def forward(self,x):
        x=self.model(x)
        return x

注:这里Flatten类自己写的原因是笔者使用的torch版本中没有展平类,因此需要自己构建,可以参考解决 ImportError: cannot import name ‘Flatten‘ from ‘torch.nn‘

创建网络模型

mxc=MXC()

创建损失函数

loss_fn=nn.CrossEntropyLoss()

设置优化器

learning_rate=1e-2
optimizer=torch.optim.SGD(params=mxc.parameters(),lr=learning_rate)

这里params是网络模型;lr是学习速率,一般设置小一些(0.01)。

训练步骤+测试步骤开始

先设置一些参数

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

将训练步骤和测试步骤放入一个大循环中,进入循环开始训练:

for i in range(epoch):
    print("-----第{}轮训练开始------".format(i+1))
    #训练步骤开始
    for data in train_data_loader:
        imgs,targets=data
        output=mxc(imgs)
        loss=loss_fn(output,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))  # 更正规的可以写成loss.item()
    #测试步骤开始
    total_test_loss=0
    with torch.no_grad():
        for data in test_data_loader:
            imgs,targets=data
            output=mxc(imgs)
            loss=loss_fn(output,targets)
            total_test_loss=total_test_loss+loss
            accuracy=(output.argmax(1)==targets).sum()
            total_accuracy=total_accuracy+accuracy
    print("整体测试集上的Loss:{}".format(total_test_loss))
    print("整体测试集上的准确率:{}".format(total_accuracy/test_data_size))
    total_test_step=total_test_step+1#测试的次数,其实就是第几轮
    #保存模型
    torch.save(mxc,"mxc_cpu{}.pth".format(i+1))#mxc_1是cpu版的
    print("模型已保存")

这里每一个data中包含有图片+标签,需要将图片(imgs)放入之前搭建好的神经网络模型mxc中去。

笔者这里设置的是每训练100次打印1次,训练完一轮后会将模型进行保存,注意文件类型是pth格式。

如果想让训练后的结果可视化,有两种方法:

1.在循环前调用SummaryWriter添加tensorboard:

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

并在循环的适当位置中插入writer.add_scalar

        if total_train_step%100==0:
            print("训练次数:{} , Loss:{}".format(total_train_step, loss))  # 更正规的可以写成loss.item()
            writer.add_scalar("train_loss",loss.item(),total_train_step)
print("整体测试集上的Loss:{}".format(total_test_loss))
print("整体测试集上的准确率:{}".format(total_accuracy/test_data_size))
total_test_step=total_test_step+1#测试的次数,其实就是第几轮
writer.add_scalar("test_loss",total_test_loss,total_test_step)
writer.add_scalar("test_accuracy",total_accuracy/test_data_size,total_test_step)

在运行结束后,打开Terminal输入(注意,前面需要显示pytorch,因为只有在pytorch环境下才可以,如果没有显示还要切换到pytorch才行,可以输入activate pytorch):

tensorboard --logdir=logs_train --port=6007

这里的“logs_train”是在SummaryWriter中设置的保存路径。运行之后,就可以在tensorboard中查看随着训练次数的增加测试集上的Loss和准确率的趋势图像。

2.调用matlab库自己进行绘制

总结

本文简要介绍了pytorch模型训练的一个基本流程,并以CIFAR10数据集进行了演示。

但这种方法是在CPU(device=“cpu”)上进行训练的,训练速度比较慢,如果数据集十分庞大不建议使用这种方法,应该在GPU(device=“cuda”)上进行训练。

以上为个人经验,希望能给大家一个参考,也希望大家多多支持脚本之家。

您可能感兴趣的文章:
阅读全文