初识Pytorch(三) -- 完整的模型训练流程 + GPU调用训练

1 构建训练模型

以CIFAR10数据集为例

1.1 导入torch模块

1
2
3
import torchvision
from torch import nn
from torch.utils.data import DataLoader

1.2 准备数据集

1
2
3
4
5
6
7
8
9
10
# 准备数据集
train_data = torchvision.datasets.CIFAR10(
    root='./dataset', train=True, download=True,
    transform=torchvision.transforms.ToTensor()
)

test_data = torchvision.datasets.CIFAR10(
    root='./dataset', train=False, download=True,
    transform=torchvision.transforms.ToTensor()
)

1.3 查看数据集的大小

1
2
3
4
5
# length 长度
train_data_size = len(train_data)
test_data_size = len(test_data)
print('训练数据集的长度为：{}'.format(train_data_size))
print('测试数据集的长度为：{}'.format(test_data_size))

1.4 加载数据集

1
2
3
# 利用dataloader加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

1.5 搭建模型

以 CIFAR 10 的结构为例，神经网络模型结构如下

image-20250604122341166

创建model.py文件存放自定义模型

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# model.py文件
import torch
from torch import nn

class TdModel(nn.Module):
    def __init__(self):
        super(TdModel, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=5, stride=1, padding=2),
            nn.MaxPool2d(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(),
            nn.Linear(64 * 4 * 4, 64),
            nn.Linear(64, 10)
        )

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

1.6 训练参数设置

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
from model import *

# 创建网络模型
td = TdModel()

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

# 优化器
# learning_rate = 0.01    # 学习率
# 1e-2 = 1*(10)^(-2) = 1/100 = 0.01
learning_rate = 1e-2    # 学习率
optimizer = torch.optim.SGD(td.parameters(), lr=learning_rate)

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

1.7 模型训练与测试

使用tensorboard记录训练过程，查看模型是否训练达到自己的需求：

每次训练完一轮后，在测试数据集上跑一遍，用测试数据集的损失或正确率来评估模型有没有训练好。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
# 添加tensorboard
writer = SummaryWriter('logs')

# 训练内容
for i in range(epoch):
    print('--------------------第{}轮训练开始--------------------'.format(i+1))
    
    # 训练步骤开始
    for data in train_dataloader:
        imgs, targets = data
        outputs = td(imgs)
        loss = lose_fn(outputs, targets)    # 损失值
        
        # 优化器模型
        optimizer.zero_grad()               # 梯度清零
        loss.backward()                     # 计算每个权重的梯度
        optimizer.step()                    # 进行权重优化
        
        total_train_step += 1               # 训练次数统计
        if total_train_step % 100 == 0:     # 训练步骤逢百打印
            print('训练次数：{0}，Loss：{1}'.format(total_train_step, loss.item()))
            writer.add_scalar('train_loss', loss.item(), total_train_step)

    # 测试步骤开始
    total_test_loss = 0
    with torch.no_grad():    # 禁用梯度计算，不去track梯度
        for data in test_dataloader:
            imgs, targets = data
            outputs = td(imgs)        # 使用已训练的模型进行推理
            loss = lose_fn(outputs, targets)
            total_test_loss += loss.item()    # loss为tensor数据类型，loss.item()为普通数字类型
        print('整体测试集上的Loss：{}'.format(total_test_loss))
        writer.add_scalar('test_loss', total_test_loss, total_test_step, )
        total_test_step += 1

        torch.save(td, './td_model{}.pt'.format(i))
        print('模型已保存')

1.8 正确率(分类问题)

即便得到整体数据集上的Loss，也不能很好地说明模型在测试集上的效果

在分类问题中可以用正确率来表示模型是否优秀

对于目标检测、语义分割等操作，可以直接把得到的输出在tensorboard中显示

正确率测试说明

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import torch

outputs = torch.tensor(         # 自定义三种类别，两个样本（一行为一个样本的数据）
    [[0.1,0.2,0.8],
     [0.5,0.05,0.4]]
)
print(outputs.argmax(0))    # tensor([1, 0, 0])
print(outputs.argmax(1))    # tensor([2, 0])
# argmax(1)水平方向，对每一行操作，返回该行中最大值的索引
# argmax(0) 垂直方向，对每一列操作，返回该列中最大值的索引

# 样本验证
preds = outputs.argmax(1)           # 找到每个样本中概率最大的索引
targets = torch.tensor([2,1])       # 只有两个样本，样本的标签对应只有两个！！！
true_number = (preds == targets)    # tensor([ True, False])
print(true_number.sum())            # 计算对应位置相等的个数
# tensor(1)

将1.7代码优化，计算测试集中整体的正确个数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# 测试步骤开始
total_test_loss = 0
total_test_accuracy = 0
with torch.no_grad():
    for data in test_dataloader:
        imgs, targets = data
        outputs = td(imgs)
        loss = lose_fn(outputs, targets)
        total_test_loss += loss.item()
        
    # 计算测试集上的正确率    
        accuracy = (outputs.argmax(dim=1) == targets).sum()
        total_test_accuracy += accuracy.item()
    print('整体测试集上的正确率：{}'.format(total_test_accuracy/test_data_size))
    writer.add_scalar('test_accuracy', total_test_accuracy/test_data_size, total_test_step)
        
    print('整体测试集上的Loss：{}'.format(total_test_loss))
    writer.add_scalar('test_loss', total_test_loss, total_test_step, )
    total_test_step += 1
    
    # 保存每轮训练的结果
    torch.save(td, './td_model_{}.pth'.format(i))
    print('第{}次训练模型已保存'.format(i))

1.9 完整代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

# 1.准备数据集
train_data = torchvision.datasets.CIFAR10(
    root='./dataset', train=True, download=True,
    transform=torchvision.transforms.ToTensor()
)
test_data = torchvision.datasets.CIFAR10(
    root='./dataset', train=False, download=True,
    transform=torchvision.transforms.ToTensor()
)
# 数据集 length 长度
train_data_size = len(train_data)
test_data_size = len(test_data)
print('训练数据集的长度为：{}'.format(train_data_size))
print('测试数据集的长度为：{}'.format(test_data_size))

# 2.准备dataloader加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

# 3.创建网络模型
from model import *    # 网络模型TdModel()在model.py文件中
td = TdModel()

# 4.损失函数
lose_fn = nn.CrossEntropyLoss()

# 5.优化器
# learning_rate = 0.01    # 学习率
# 1e-2 = 1*(10)^(-2) = 1/100 = 0.01
learning_rate = 1e-2    # 学习率
optimizer = torch.optim.SGD(td.parameters(), lr=learning_rate)

# 6.设置训练网络的一些参数
total_train_step = 0    # 记录训练的次数
total_test_step = 0     # 记录测试的次数
epoch = 50              # 训练的轮次 

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

# 7.开始训练内容，以每个epoch为一轮回，实现多次训练
for i in range(epoch):
    print('--------------------第{}轮训练开始--------------------'.format(i+1))
    
    # 8.训练步骤开始，模型网络进入训练状态
    # td.train()    # 这仅对某些模块有效，例如 Dropout、BatchNorm 等
    for data in train_dataloader:    # 从训练的dataloader中不断的取数据
        imgs, targets = data
        outputs = td(imgs)
        loss = lose_fn(outputs, targets)    # 计算损失值
        
        # 9.优化器模型，计算出误差后放入优化器中进行优化
        optimizer.zero_grad()               # 梯度清零
        loss.backward()                     # 计算每个权重的梯度
        optimizer.step()                    # 进行权重优化
        
        # 10.展示输出
        total_train_step += 1               # 训练次数统计
        if total_train_step % 100 == 0:
            print('训练次数：{0}，Loss：{1}'.format(total_train_step, loss.item()))
            writer.add_scalar('train_loss', loss.item(), total_train_step)
    
    # 11.测试步骤开始，在 一轮结束 或 特定步数 后进行测试
    # td.eval()    # 这仅对某些模块有效。例如 Dropout、BatchNorm 等
    total_test_loss = 0
    total_test_accuracy = 0
    
    with torch.no_grad():          # 禁用梯度计算，不去track梯度
        for data in test_dataloader:
            imgs, targets = data
            outputs = td(imgs)
            loss = lose_fn(outputs, targets)
            total_test_loss += loss.item()
            
        # 计算测试集上的正确率    
            accuracy = (outputs.argmax(dim=1) == targets).sum()
            total_test_accuracy += accuracy.item()
        print('整体测试集上的正确率：{}'.format(total_test_accuracy/test_data_size))
        writer.add_scalar('test_accuracy', total_test_accuracy/test_data_size, total_test_step)
            
        print('整体测试集上的Loss：{}'.format(total_test_loss))
        writer.add_scalar('test_loss', total_test_loss, total_test_step, )
        total_test_step += 1
        
        # 保存每轮训练的结果
        torch.save(td, './td_model_{}.pth'.format(i))
        print('第{}次训练模型已保存'.format(i))

2 使用 GPU 训练

2.1 调用方式 1

只有网络模型、损失函数、数据（输入、输出）可以设置为GPU运算

使用方式：在对象后添加 .cuda()

网络模型

1
2
3
4
5
# 3.创建网络模型
from model import *    # 网络模型TdModel()在model.py文件中
td = TdModel()
if torch.cuda.is_available():
    td = td.cuda()

损失函数

1
2
3
4
# 4.损失函数
lose_fn = nn.CrossEntropyLoss()
if torch.cuda.is_available():
    lose_fn = lose_fn.cuda()

输入、输出数据

1
2
3
4
5
6
7
8
# 8.训练步骤开始，模型网络进入训练状态
# td.train()    # 这仅对某些模块有效，例如 Dropout、BatchNorm 等
for data in train_dataloader:    # 从训练的dataloader中不断的取数据
    imgs, targets = data
    if torch.cuda.is_available():
        imgs, targets = imgs.cuda(), targets.cuda()
    outputs = td(imgs)
    loss = lose_fn(outputs, targets)    # 计算损失值

2.2 调用方式 2

使用方式：在对象后添加 .to(device)

1
2
3
4
5
6
device = torch.device('cpu')    # 指定运算设备为cpu

device = torch.device('cuda')      # 指定运算设备为gpu，默认为第一张显卡
device = torch.device('cuda:0')    # 指定运算设备为gpu，使用第一张显卡

device = torch.device('cuda:1')    # 指定运算设备为gpu，使用第二张显卡

指定设备

1
2
# 指定训练的设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

调用设备

1
2
3
4
5
6
7
8
9
10
11
12
13
# 3.创建网络模型
from model import *    # 网络模型TdModel()在model.py文件中
td = TdModel()
td = td.to(device)

# 4.损失函数
lose_fn = nn.CrossEntropyLoss()
lose_fn = lose_fn.to(device)

# 从训练的dataloader中取数据
for data in train_dataloader:    
    imgs, targets = data
    imgs, targets = imgs.to(device), targets.to(device)

3 测试模型结果

利用已经训练好的模型，给模型提供输入

从网络上下载一个图片，使用PIL读取图片

因为png格式是四个通道，除了RGB三通道外，还有一个透明度通道

调用 image = image.convert(‘RGB’) 保留图片的颜色通道

如果图片本来就是三个颜色通道，经过此操作后，无变化

加上这一步后可以适应png、jpg等格式的图片

1
2
3
4
5
6
7
8
9
from PIL import Image
import torchvision

# 使用PIL读取图片
image_path = 'images/dog.png'
image = Image.open(image_path)  
image = image.convert('RGB')
print(image)
# <PIL.Image.Image image mode=RGB size=197x143 at 0x1991212BE20>

调整图像大小

1
2
3
4
5
6
transform = torchvision.transforms.Compose([    # 联立transform
    torchvision.transforms.Resize((32,32)),     # 对输入图像进行尺寸调整。
    torchvision.transforms.ToTensor(),
])
image = transform(image)
print(image.shape)

加载网络模型

1
2
3
4
5
# 加载网络模型
from model import *
model = torch.load("td_model_0.pth",map_location=torch.device('cpu'), weights_only=False)
print(model)
# 注意此时引用的模型文件是用gpu训练的，如果测试时使用的是cpu，则需要指定设备到cpu去运算

输入图片，推理结果

1
2
3
4
5
6
image = torch.reshape(image, (1,3,32,32))
model.eval()            # 用于将模型设置为评估模式。
with torch.no_grad():   # 禁用梯度计算（节省内存，加速推理）。
    output = model(image)
print(output)
print(output.argmax(dim=1))