Pytorch 实战四 VGG 网络训练

系列文章目录

文章目录

系列文章目录
前言
一、源码
- [1. 解决线程冲突](#1. 解决线程冲突)
- 2.代码框架
二、代码详细介绍
- 1.基础定义
- [2. epoch 的定义](#2. epoch 的定义)
- [3. 每组图片的训练和模型保存](#3. 每组图片的训练和模型保存)

前言

前面我们已经完成了数据集的制作，VGG 网络的搭建，现在进行网络模型的训练。

一、源码

python 复制代码

import torch.nn as nn
import torchvision
from VggNet import VGGNet
from load_cifa10 import train_data_loader, test_data_loader
import torch.multiprocessing as mp
import torch
import multiprocessing
from torch.utils.data import DataLoader

from model.ClassModel import net


def main():
    # 训练模型到底放在 CPU 还是GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 如果cuda有效，就在GPU训练，否则CPU训练

    # 我们会对样本遍历20次
    epoch_num = 20
    # 学习率
    lr = 0.01
    # 正确率计算相关
    batch_num=0
    correct0 =0

    # 网络定义
    # print("need初始化")
    net = VGGNet().to(device)

    # 定义损失函数loss,多分类问题，采用交叉熵
    loss_func = nn.CrossEntropyLoss()
    # 定义优化器optimizer
    optimizer = torch.optim.Adam(net.parameters(),lr=lr)

    # 动态调整学习率,第一个参数是优化器，起二个参数每5个epoch后调整学习率，第三个参数调整为原来的0.9倍
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5,gamma=0.9)

    # 定义循环
    for epoch in range(epoch_num):
        # print("epoch is:",epoch+1)
        # 定义网络训练的过程
        net.train()   # BatchNorm 和 dropout 会选择相应的参数
        # 对数据进行遍历
        for i,data in enumerate(train_data_loader):
            batch_num = len(train_data_loader)
            # 获取输入和标签
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device) # 放到GPU上去

            # 拿到输出
            # print("need output")
            outputs = net(inputs)  # 这句就会调用前向传播，在PyTorch中，当执行outputs = net(inputs)时会自动触发前向传播，
                                    # 这是通过nn.Module的__call__方法实现的特殊机制6。具体原理可分为三个关键环节：
            # 计算损失
            loss = loss_func(outputs, labels)
            # 定义优化器，梯度要归零，loss反向传播，更新参数
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()


            _, predicted = torch.max(outputs.data, dim=1)  # 得到一个batch的预测类别
            # 在cpu上面运行,当labels是普通张量时，.data属性返回‌剥离计算图的纯数值张量‌（与原始张量共享内存但无梯度追踪）
            correct = predicted.eq(labels.data).cpu().sum()
            correct=100.0*correct/len(inputs)
            correct0+=correct
        #自动更新学习率
        lr_scheduler.step()
        lr = optimizer.state_dict()['param_groups'][0]['lr']
    print("loss:{},acc:{}",lr,correct0/batch_num)

if __name__ == '__main__':
    # Windows必须设置spawn，Linux/Mac自动选择最佳方式
    mp.set_start_method('spawn' if torch.cuda.is_available() else 'fork')
    torch.multiprocessing.freeze_support()

    try:
        main()
    except RuntimeError as e:
        print(f"多进程错误: {str(e)}")
        print("降级到单进程模式...")
        train_data_loader = DataLoader(..., num_workers=0)
        main()

1. 解决线程冲突

windows 跑代码需要解决线程冲突的问题：需要自行定义main函数，然后把主题加在里面。当我们运行时自动调用main，就会执行下面的 if 语句，然后运行我们的代码

python 复制代码

if __name__ == '__main__':
    # Windows必须设置spawn，Linux/Mac自动选择最佳方式
    mp.set_start_method('spawn' if torch.cuda.is_available() else 'fork')
    torch.multiprocessing.freeze_support()

    try:
        main()
    except RuntimeError as e:
        print(f"多进程错误: {str(e)}")
        print("降级到单进程模式...")
        train_data_loader = DataLoader(..., num_workers=0)
        main()

2.代码框架

代码分成四个部分，第一个部分是基础变量定义，第二个部分是循环 epoch ,第三部分是每个 batch 的处理，第四个保存模型，其中最重要的便是第三个。
图 1 代码框架

二、代码详细介绍

1.基础定义

python 复制代码

    # 训练模型到底放在 CPU 还是GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 如果cuda有效，就在GPU训练，否则CPU训练

    # 我们会对样本遍历200次
    epoch_num = 200
    # 学习率
    lr = 0.01
    # 正确率计算相关
    batch_num=0
    correct0 =0

    # 网络定义
    # print("need初始化")
    net = VGGNet().to(device)

    # 定义损失函数loss,多分类问题，采用交叉熵
    loss_func = nn.CrossEntropyLoss()
    # 定义优化器optimizer
    optimizer = torch.optim.Adam(net.parameters(),lr=lr)

    # 动态调整学习率,第一个参数是优化器，起二个参数每5个epoch后调整学习率，第三个参数调整为原来的0.9倍
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5,gamma=0.9)

基础定义最开始需要定义跑数据的设备，CPU还是GPU.这个是死的。然后定义epoch次数、学习率，至于准确率看自己的使用情况，如果每次跑完一遍数据集打印准确率也不着急定义。我在最终跑完数据才打印准确率，所以需要定义一个全局的变量。接下来便是网络初始化，初始化的网络加载到设备上面。在网络搭建的时候，我们只定义了网络的层次和前向传播。后面的损失函数和优化器需要在训练中进行。那么基础定义里面需要损失函数的选择，优化器的选择和动态调整学习率。epoch 改成200，我的电脑跑了1h还没出结果，现在还在等，建议别弄大了。
图 2 基础定义

当然顺序可以变，最好自己能记住需要的内容。

2. epoch 的定义

epoch里面开始调用网络，net.train() 会把网络的参数进行初始化，BatchNorm 会自动启用训练模式，dropout层会全部激活，而这在测试集上不需要dropout的。后面便是每组图片 batch的训练行为。最后每一次处理整个数据集需要动态改变学习率，以及打印学习率的方法如下：

python 复制代码

        #自动更新学习率
        lr_scheduler.step()
        # 打印学习率
        lr = optimizer.state_dict()['param_groups'][0]['lr']
        print("学习率", lr)

3. 每组图片的训练和模型保存

python 复制代码

      for i,data in enumerate(train_data_loader):
            batch_num = len(train_data_loader)
            # 获取输入和标签
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device) # 放到GPU上去

            # 拿到输出
            # print("need output")
            outputs = net(inputs)  # 这句就会调用前向传播，在PyTorch中，当执行outputs = net(inputs)时会自动触发前向传播，
                                    # 这是通过nn.Module的__call__方法实现的特殊机制6。具体原理可分为三个关键环节：
            # 计算损失
            loss = loss_func(outputs, labels)
            # 定义优化器，梯度要归零，loss反向传播，更新参数
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()


            _, predicted = torch.max(outputs.data, dim=1)  # 得到一个batch的预测类别

这里进行数据加载，分批次加载，此处的batch 大小是128，数量是391（用于准确率计算）。加载了数据，获取数据的输入和真实标签。outputs = net(inputs) 这句对网络传入数据，自行前向传播计算获得输出。拿到输出后，进行损失函数计算。损失函数计算，是需要预测值和真实值的，看看偏差多少，因此传入这两个参数。优化器优化，开始梯度归零，然后后向传播，这个过程是自带的，我们只定义了前向传播，后向传播优化参数后固定参数 optimizer.step(),最后使用torch.max（） 输出最相似的标签。过程如图：
图 3 batch 循环

模型的保存就一句话：

python 复制代码

torch.save(net.state_dict(),"./model/VGGNet.pth")