利用Pytorch搭建简单的图像分类模型（之三）---训练与测试的函数编写

不换心

已于 2022-09-27 21:11:17 修改

阅读量706

点赞数 1

文章标签：深度学习 pytorch python

于 2022-09-27 20:45:38 首次发布

本文链接：https://blog.csdn.net/Hust_GreenHand/article/details/127078248

版权

利用Pytorch来实现模型的训练和测试

1.每一个epoch的训练步骤

epoch的意思是完整训练模型的次数，注意与batchsize区分开。训练的步骤无非就是传入数据 → 前向传播 → 计算损失 → 反向传播 → 更新权重，那么下面就按照这个步骤来编写一个epoch的训练步骤。

import sys
import torch
import torch.nn as nn
from tqdm import tqdm

def train_each_epoch(model, optimizer, train_dataloader, device, epoch):
    model.train()   # 设置为训练模式，训练的时候一定要加这一句，这样才会计算并更新BatchNormalization和采用Dropout

    '''初始化参数和指定损失函数'''
    loss_function = nn.CrossEntropyLoss()   # 分类任务最常用的就是交叉熵损失函数
    loss_sum = torch.zeros(1).to(device)    # 初始化总损失，得到一个初值为0的1*1的tensor
    acc_cnt = torch.zeros(1).to(device)     # 初始化预测正确的样本数量
    sample_cnt = 0  # 初始化总样本数量，以便后续求解准确率等指标，这里的初始化不用torch.zeros(1).to(device)这种形式是因为sample_cnt不会用到模型中输出的结果
    optimizer.zero_grad()   # 将初始梯度清零，避免干扰训练结果

    '''训练步骤'''
    train_dataloader = tqdm(train_dataloader)   # 引入tqdm，tqdm是一个可扩展的Python进度条，可以在Python长循环中添加一个进度提示信息，也就是可视化我们的训练过程
    for step, data in enumerate(train_dataloader):  # 一个batchsize就是一个step
        images, labels = data   # 得到一个batchsize的图片数据和其对应标签
        
        '''统计样本数量'''
        sample_cnt += images.shape[0]    # 一个step的image为（64, 3, 224, 224）的tensor，shape[0]返回的就是最外围的维度（64）就是一个batchsize中包含样本的数量

        '''统计预测正确的数量'''
        pred = model(images.to(device))  # 这里得到的预测结果是一个(64, 5)的tensor，5代表数据集类别的数量，这里的输出只是每个类别的概率
        pred_classes = torch.max(pred, dim=1)[1]    # 在5个概率里面选择最大的那个，并返回最大值的索引，[ ]中是0的话就是返回最大的值，是1的话就是返回最大值的索引
        acc_cnt = torch.eq(pred_classes, labels.to(device)).sum()   # 统计一个batchsize中预测正确的数量torch.eq()就是判断是否相等

        '''计算损失，并反向传播'''
        loss = loss_function(pred, labels.to(device))   # 将预测的概率与标签去求损失
        loss.requires_grad_(True)   # 作用是让backward可以追踪这个参数并且计算它的梯度，训练的时候一定要把损失的requires_grad_()设置为True，不然无法反向传播
        loss.backward()     # 开始反向传播，计算每一个参与运算的权重W的梯度
        '''将每一个step的损失相加得到总损失'''
        loss_sum += loss.detach()   # 这里记得一定要加.detach()或者loss.data也可以，因为loss_sum是不需要参与梯度计算的，
                                    # 加.detach()是为了阻断后续Pytorch继续构图，不加的话随着step的增加，显存会占用的越来越多。

        '''调用tqdm中的.desc,可视化训练过程'''
        """
            epoch: 显示训练多少个epoch了
            loss_sum.item()/(step+1): 将总损失除以当前迭代了多少个step,得到当前每个step的平均损失
            acc_cnt.item()/sample_cnt: 计算当前预测的准确率
        """
        train_dataloader.desc = "[train epoch {}] loss: {:3f}, acc: {:3f}".format(
            epoch, loss_sum.item()/(step+1), acc_cnt.item()/sample_cnt
        )

        '''这里是为了检验有没有出现不合法的梯度,如果有的话直接终止训练(意味着要去调Bug了!)'''
        if not torch.isfinite(loss):
            print("WARNING: non-finite loss, ending training ", loss)
            sys.exit(1)

        optimizer.step()    # 开始根据优化器算法进行权重更新，一定要放在loss.backward()后面
        optimizer.zero_grad()   # 将梯度清零，进行下一个step的训练

    return  loss_sum.item()/len(train_dataloader), acc_cnt.item()/sample_cnt    # 完成一个epoch的训练后，返回这个epoch的平均损失和正确率

上述代码就是在模型训练时候，每一个epoch的训练步骤，下面就来看看测试代码怎么写，其实了解了训练的步骤，测试步骤就很简单了，只有几个地方需要修改一下。

2.测试代码的编写

import torch

def test(model, test_dataloader, device):
    model.eval()

    acc_cnt = torch.zeros(1).to(device)
    sample_cnt = 0
    with torch.no_grad():
        for test_data in tqdm(test_dataloader):

            test_images, test_labels = test_data
            sample_cnt += test_images.shape[0]

            pred = model(test_images.to(device))
            pred_classes = torch.max(pred, dim=1)[1]

            acc_cnt += torch.eq(pred_classes, test_labels.to(device)).sum()
    
    print("The test Acc is: {:3f}".format(acc_cnt.item()/sample_cnt))