完整的模型训练套路：代码模板

数据集以经典的 CIFAR10 为例。

这个例子是很简单的，可能不太实用，但重点是通过这个例子掌握一种模型训练的写法套路，因此很有必要学习。

import torch.optim
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="./dataset2", train=True,transform=torchvision.transforms.ToTensor(), download=True)# 2、准备测试数据集
test_data = torchvision.datasets.CIFAR10(root="./dataset2", train=False,transform=torchvision.transforms.ToTensor(), download=True)# 3、查看数据集的大小
train_data_size = len(train_data)
test_data_size = len(test_data)
print("训练数据集的长度为：{}".format(train_data_size))
print("测试数据集的长度为：{}".format(test_data_size))# 4、利用 DataLoader 加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)# 5、搭建神经网络：完成一个十分类的任务
class NetWork(nn.Module):def __init__(self):super(NetWork, 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),nn.Flatten(),nn.Linear(64 * 4 * 4, 64),nn.Linear(64, 10),)def forward(self, x):x = self.model(x)return x# 6、创建网络模型
network = NetWork()# 7、定义损失函数
loss_fn = nn.CrossEntropyLoss()# 8、定义优化器
# 为神经网络network定义一个SGD（随机梯度下降）优化器，并使用指定的学习率learning_rate来更新其参数。
learning_rate = 0.01
# torch.optim.SGD：这是PyTorch中SGD优化器的定义。SGD是一种常用的优化算法，用于在训练过程中最小化损失函数。
# network.parameters()：这是一个生成器，它返回神经网络network中所有可训练的参数（如权重和偏置）。
#                       在PyTorch中，模型通常是一个nn.Module的子类，并且模型中的参数可以通过.parameters()方法获取。
# lr=learning_rate：这是SGD优化器的一个参数，表示学习率（learning rate）。
#                   学习率是一个超参数，用于控制参数更新的步长。较大的学习率可能导致训练不稳定，而较小的学习率可能导致训练速度较慢。
#                   你需要根据具体的任务和模型结构来选择合适的学习率。
optimizer = torch.optim.SGD(network.parameters(), lr=learning_rate)# 9、设置训练网络的一些参数
# 记录训练的次数
total_train_step = 0
# 记录测试的次数
total_test_step = 0
# 训练的轮数
epoch = 10# 10、添加 tensorboard 显示
writer = SummaryWriter("./logs_train")# 11、开始训练
for i in range(epoch):print("----------第 {} 轮训练开始----------".format(i+1))# 训练步骤开始# network.train() 这一步对于当前网络来说并不是必要的，有无皆可，但如果网络中涉及了 Dropout 层 或者 BatchNorm 层等才必须加上for data in train_dataloader:imgs, targets = dataoutputs = network(imgs)# 计算模型预测与真实目标之间的损失# 它接受模型的输出（outputs）和真实的目标值（targets）作为输入，并返回一个表示两者差异的数值（通常是一个标量张量）。loss = loss_fn(outputs, targets)# 优化器开始调优，优化模型参数# 在每次训练迭代（iteration）开始时，都会调用这个zero_grad()方法来清除之前累积的梯度。# 这是因为在PyTorch中，梯度是累积的，如果你不清除之前的梯度，# 那么在下一次迭代时，梯度就会是之前迭代梯度的累加，这通常不是我们想要的。optimizer.zero_grad()# backward()这个方法会计算损失函数关于模型参数的梯度，并将结果存储在参数的.grad属性中。这是反向传播（backpropagation）的核心步骤。loss.backward()# step()这个方法会根据之前计算得到的梯度来更新模型的参数。# 具体来说，它会根据优化器（如SGD）中定义的算法来更新模型的参数，使其朝着损失函数减小的方向前进。optimizer.step()total_train_step = total_train_step + 1# 控制一下打印的频率，让其每逢100次时再打印日志if total_train_step % 100 == 0:# 在 PyTorch 中，当你计算得到一个损失（loss）或任何单元素张量（scalar tensor）时，这个值通常是一个 PyTorch 张量（tensor）。# 然而，在很多情况下，你可能想要将这个张量值转换为一个 Python 的基本数据类型（如 float 或 int），以便进行更简单的数学运算或打印。# loss.item() 就是用来做这个转换的。它会返回张量中的单个元素值作为一个 Python 数字。# 注意，这个方法只能在单元素张量（scalar tensor）上调用，即那些形状为 torch.Size([]) 的张量。print("训练次数: {}, Loss: {}".format(total_train_step, loss.item()))# 在 TensorBoard 中添加一个名为 'train_loss' 的标量（scalar）数据点，# 其值为当前的损失值（loss.item()），并将其与训练步骤 total_train_step 相关联”# "train_loss"：这是你要在 TensorBoard 中显示的数据的标签。#               当你打开 TensorBoard 并查看数据时，你会在图表或表格中看到这个标签。# loss.item()：这是从 PyTorch 的张量（tensor）中提取出的单个元素值（在这个情况下是损失值）。# total_train_step：这是训练步骤的编号或索引。在训练过程中，随着每个批次的迭代，这个值通常会递增。#                   它用于在 TensorBoard 的图表中跟踪损失值随时间（或训练步骤）的变化。writer.add_scalar("train_loss", loss.item(), total_train_step)# 测试步骤开始（每完成一个轮回，即一次前向和反向传播完整个训练数据集，就测试一下效果，看测试集上的损失值大小）# network.eval() 这一步对于当前网络来说并不是必要的，有无皆可，但如果网络中涉及了 Dropout 层 或者 BatchNorm 层等才必须加上total_test_loss = 0# 在 PyTorch 中，torch.no_grad() 并不是一个函数调用，而是一个上下文管理器（context manager），# 通常与 with 关键字一起使用，以确保在 with 语句块中的操作不会计算梯度，从而节省内存和计算资源。with torch.no_grad():for data in test_dataloader:imgs, targets = dataoutputs = network(imgs)loss = loss_fn(outputs, targets)total_test_loss = total_test_loss + loss.item()print("整体测试集上的Loss: {}".format(total_test_loss))writer.add_scalar("test_loss", total_test_loss, total_test_step)total_test_step = total_test_step + 1# 12、保存下每一轮训练的模型torch.save(network.load_state_dict(), "network_{}.pth".format(i+1))print("模型已保存")# 13、关闭 tensorboard
writer.close()

上述代码就是目标检测或者分割类 CV 任务的基本模板了，但其实我们看很多的开源项目或者是别人的代码，习惯上都喜欢将第五步搭建神经网络部分给单独封装到一个 python 文件当中从而进行调用，因此我们也可以仿照着这样做。

将搭建神经网络的部分单独封装成一个 model.py 文件（注意要和我们自己的训练文件位于同一目录下嗷）：

# 5、搭建神经网络：完成一个十分类的任务
import torch
from torch import nn# 5、搭建神经网络：完成一个十分类的任务
class NetWork(nn.Module):def __init__(self):super(NetWork, 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),nn.Flatten(),nn.Linear(64 * 4 * 4, 64),nn.Linear(64, 10),)def forward(self, x):x = self.model(x)return x# 在这里测试一下网络的正确性
if __name__ == '__main__':network = NetWork()# 看一下输出的尺寸是不是我们想要的，如果是那么就应该是正确的input = torch.ones((64, 3, 32, 32))output = network(input)print(output.shape)

然后在我们的训练文件 train.py 中进行引入：

import torch.optim
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from model import *# 1、准备训练数据集
train_data = torchvision.datasets.CIFAR10(root="./dataset2", train=True,transform=torchvision.transforms.ToTensor(), download=True)# 2、准备测试数据集
test_data = torchvision.datasets.CIFAR10(root="./dataset2", train=False,transform=torchvision.transforms.ToTensor(), download=True)# 3、查看数据集的大小
train_data_size = len(train_data)
test_data_size = len(test_data)
print("训练数据集的长度为：{}".format(train_data_size))
print("测试数据集的长度为：{}".format(test_data_size))# 4、利用 DataLoader 加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)# 5、引入我们在 model.py 中搭建好的神经网络：完成一个十分类的任务# 6、创建网络模型
network = NetWork()# 7、定义损失函数
loss_fn = nn.CrossEntropyLoss()# 8、定义优化器
# 为神经网络network定义一个SGD（随机梯度下降）优化器，并使用指定的学习率learning_rate来更新其参数。
learning_rate = 0.01
# torch.optim.SGD：这是PyTorch中SGD优化器的定义。SGD是一种常用的优化算法，用于在训练过程中最小化损失函数。
# network.parameters()：这是一个生成器，它返回神经网络network中所有可训练的参数（如权重和偏置）。
#                       在PyTorch中，模型通常是一个nn.Module的子类，并且模型中的参数可以通过.parameters()方法获取。
# lr=learning_rate：这是SGD优化器的一个参数，表示学习率（learning rate）。
#                   学习率是一个超参数，用于控制参数更新的步长。较大的学习率可能导致训练不稳定，而较小的学习率可能导致训练速度较慢。
#                   你需要根据具体的任务和模型结构来选择合适的学习率。
optimizer = torch.optim.SGD(network.parameters(), lr=learning_rate)# 9、设置训练网络的一些参数
# 记录训练的次数
total_train_step = 0
# 记录测试的次数
total_test_step = 0
# 训练的轮数
epoch = 10# 10、添加 tensorboard 显示
writer = SummaryWriter("./logs_train")# 11、开始训练
for i in range(epoch):print("----------第 {} 轮训练开始----------".format(i+1))# 训练步骤开始# network.train() 这一步对于当前网络来说并不是必要的，有无皆可，但如果网络中涉及了 Dropout 层 或者 BatchNorm 层等才必须加上for data in train_dataloader:imgs, targets = dataoutputs = network(imgs)# 计算模型预测与真实目标之间的损失# 它接受模型的输出（outputs）和真实的目标值（targets）作为输入，并返回一个表示两者差异的数值（通常是一个标量张量）。loss = loss_fn(outputs, targets)# 优化器开始调优，优化模型参数# 在每次训练迭代（iteration）开始时，都会调用这个zero_grad()方法来清除之前累积的梯度。# 这是因为在PyTorch中，梯度是累积的，如果你不清除之前的梯度，# 那么在下一次迭代时，梯度就会是之前迭代梯度的累加，这通常不是我们想要的。optimizer.zero_grad()# backward()这个方法会计算损失函数关于模型参数的梯度，并将结果存储在参数的.grad属性中。这是反向传播（backpropagation）的核心步骤。loss.backward()# step()这个方法会根据之前计算得到的梯度来更新模型的参数。# 具体来说，它会根据优化器（如SGD）中定义的算法来更新模型的参数，使其朝着损失函数减小的方向前进。optimizer.step()total_train_step = total_train_step + 1# 控制一下打印的频率，让其每逢100次时再打印日志if total_train_step % 100 == 0:# 在 PyTorch 中，当你计算得到一个损失（loss）或任何单元素张量（scalar tensor）时，这个值通常是一个 PyTorch 张量（tensor）。# 然而，在很多情况下，你可能想要将这个张量值转换为一个 Python 的基本数据类型（如 float 或 int），以便进行更简单的数学运算或打印。# loss.item() 就是用来做这个转换的。它会返回张量中的单个元素值作为一个 Python 数字。# 注意，这个方法只能在单元素张量（scalar tensor）上调用，即那些形状为 torch.Size([]) 的张量。print("训练次数: {}, Loss: {}".format(total_train_step, loss.item()))# 在 TensorBoard 中添加一个名为 'train_loss' 的标量（scalar）数据点，# 其值为当前的损失值（loss.item()），并将其与训练步骤 total_train_step 相关联”# "train_loss"：这是你要在 TensorBoard 中显示的数据的标签。#               当你打开 TensorBoard 并查看数据时，你会在图表或表格中看到这个标签。# loss.item()：这是从 PyTorch 的张量（tensor）中提取出的单个元素值（在这个情况下是损失值）。# total_train_step：这是训练步骤的编号或索引。在训练过程中，随着每个批次的迭代，这个值通常会递增。#                   它用于在 TensorBoard 的图表中跟踪损失值随时间（或训练步骤）的变化。writer.add_scalar("train_loss", loss.item(), total_train_step)# 测试步骤开始（每完成一个轮回，即一次前向和反向传播完整个训练数据集，就测试一下效果，看测试集上的损失值大小）# network.eval() 这一步对于当前网络来说并不是必要的，有无皆可，但如果网络中涉及了 Dropout 层 或者 BatchNorm 层等才必须加上total_test_loss = 0# 在 PyTorch 中，torch.no_grad() 并不是一个函数调用，而是一个上下文管理器（context manager），# 通常与 with 关键字一起使用，以确保在 with 语句块中的操作不会计算梯度，从而节省内存和计算资源。with torch.no_grad():for data in test_dataloader:imgs, targets = dataoutputs = network(imgs)loss = loss_fn(outputs, targets)total_test_loss = total_test_loss + loss.item()print("整体测试集上的Loss: {}".format(total_test_loss))writer.add_scalar("test_loss", total_test_loss, total_test_step)total_test_step = total_test_step + 1# 12、保存一下每一轮训练的模型torch.save(network.state_dict(), "network_{}.pth".format(i+1))print("模型已保存")# 13、关闭 tensorboard
writer.close()

使用GPU进行模型训练

CUDA是什么

在使用 GPU 之前，先了解一下什么是 CUDA：

CUDA（Compute Unified Device Architecture）是由NVIDIA公司推出的一种通用并行计算架构，主要用于解决复杂的计算问题。这个架构包含了CUDA指令集架构（ISA）以及GPU内部的并行计算引擎，使得GPU能够像CPU一样执行复杂的计算任务。

CUDA的主要特点包括：

并行计算：CUDA允许开发者利用GPU的并行处理能力，将计算任务划分为许多小的、可以独立执行的部分，并在多个处理器核心上同时执行。这种并行处理方式可以显著提高计算速度。
编程模型：CUDA提供了一组扩展的C、C++和Fortran语言的编程接口，使得开发者能够在GPU上编写程序。这些程序通常包括在CPU上运行的主机代码和在GPU上运行的设备代码。
应用领域：CUDA在科学计算、人工智能、深度学习、图像处理等领域得到了广泛的应用。通过在GPU上运行计算密集型任务，CUDA可以显著加速应用程序的执行。
内存管理：CUDA提供了对GPU内存的直接访问和管理，允许开发者控制数据在主机和GPU之间的传输。
架构优化：CUDA平台设计用于充分利用NVIDIA GPU的架构，包括它们的多核处理能力和高速内存访问，适用于需要线程间通信或数据重用的场景。

CUDA的出现使得开发者能够更加方便地利用GPU的计算能力，实现了用更加廉价的设备资源实现更高效的并行计算的目标。随着GPU技术的不断发展，CUDA的应用范围也在不断扩大，成为了计算领域的重要工具之一。

使用GPU进行模型训练的第一种方式（不推荐）

使用 GPU 训练模型有两种方式，方式一主要有下面这几步：

注意下面代码的改动依然是建立在上面的模板代码中的，因此我只给出了被改动地方的代码。

1、找到我们的网络模型，使用网络模型对象调用 cuda() 方法。

# 6、创建网络模型
network = NetWork()
# 调用 GPU 进行训练
if torch.cuda.is_available():network = network.cuda() # 对于网络来说，可以不进行赋值，直接network.cuda()也可

2、找到我们的损失函数定义的位置，使用损失函数对象调用 cuda() 方法

# 7、定义损失函数
loss_fn = nn.CrossEntropyLoss()
# 损失函数也要调用 GPU 进行训练
if torch.cuda.is_available():loss_fn = loss_fn.cuda() # 对于损失函数来说，可以不进行赋值，直接loss_fn.cuda()也可

3、找到我们训练数据的地方，使用的数据对象也需要调用 cuda() 方法

# 训练步骤开始# network.train() 这一步对于当前网络来说并不是必要的，有无皆可，但如果网络中涉及了 Dropout 层 或者 BatchNorm 层等才必须加上for data in train_dataloader:imgs, targets = data# 数据也要进行 GPU 调度if torch.cuda.is_available():imgs = imgs.cuda() # 对于数据来说，必须进行赋值，不能直接imgs.cuda()targets = targets.cuda() # 对于数据来说，必须进行赋值，不能直接targets.cuda()outputs = network(imgs)

还有测试数据的地方也要嗷：

    with torch.no_grad():for data in test_dataloader:imgs, targets = dataif torch.cuda.is_available():imgs = imgs.cuda() # 对于数据来说，必须进行赋值，不能直接imgs.cuda()targets = targets.cuda() # 对于数据来说，必须进行赋值，不能直接targets.cuda()outputs = network(imgs)loss = loss_fn(outputs, targets)total_test_loss = total_test_loss + loss.item()print("整体测试集上的Loss: {}".format(total_test_loss))

然后就可以运行程序了，可以和普通的 CPU 版本比较一下，甚至都不需要使用计时工具就能发现使用 GPU 明显要快很多，这里就不演示了。

使用GPU进行模型训练的第二种方式（推荐）

1、先定义一个训练设备的对象

# 定义训练的设备
# device = torch.device("cpu")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 1、准备训练数据集
train_data = torchvision.datasets.CIFAR10(root="./dataset2", train=True,transform=torchvision.transforms.ToTensor(), download=True)# 2、准备测试数据集
test_data = torchvision.datasets.CIFAR10(root="./dataset2", train=False,transform=torchvision.transforms.ToTensor(), download=True)

2、剩下的操作就和第一种方式差不多了，把下面四个部分的代码对应的修改一下，也就是将 cuda() 方法换成了 to() 方法而已

# 6、创建网络模型
network = NetWork()
network = network.to(device)  # 对于网络来说，可以不进行赋值，直接network.to()也可

# 7、定义损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.to(device) # 对于损失函数来说，可以不进行赋值，直接loss_fn.to()也可

# 训练步骤开始# network.train() 这一步对于当前网络来说并不是必要的，有无皆可，但如果网络中涉及了 Dropout 层 或者 BatchNorm 层等才必须加上for data in train_dataloader:imgs, targets = dataimgs = imgs.to(device) # 对于数据来说，必须进行赋值，不能直接imgs.to()targets = targets.to(device) # 对于数据来说，必须进行赋值，不能直接targets.to()outputs = network(imgs)

# 测试步骤开始（每完成一个轮回，即一次前向和反向传播完整个训练数据集，就测试一下效果，看测试集上的损失值大小）total_test_loss = 0with torch.no_grad():for data in test_dataloader:imgs, targets = dataimgs = imgs.to(device) # 对于数据来说，必须进行赋值，不能直接imgs.cuda()targets = targets.to(device) # 对于数据来说，必须进行赋值，不能直接targets.cuda()

运行效果就不再赘述了，可以自己测试看看，比 CPU 快不少捏。

完整的模型验证（测试）套路

模型验证的含义：使用已经训练好的模型，然后给它提供新的输入数据（测试集）进行测试。

参考 Github 上众多优秀的开源项目不难发现，在每一个项目中都有会训练文件，比如本文开头讲的完整的模型训练套路 train.py ，同理也都会有一个测试模型验证的文件 test.py ，而这一节要说的就是如何写一份完整的模型验证的套路模板。

import torch
import torchvision
from PIL import Image
from model import NetWork# 1、加载测试数据集（这里以一个简单的图片为例）
image_path = "./images/img.png"# 注意：PNG格式的图片是四通道的，记得转成三通道的进行使用，否则会报错
image = Image.open(image_path).convert("RGB")
print(image)# 2、对测试用的数据集进行转换，转换成 PyTorch 能接受的形式
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),torchvision.transforms.ToTensor(),
])image = transform(image)
print(image.shape)# 3、加载我们训练好的模型
# 注意：如果是用GPU训练出来的模型要放在CPU上跑，那么要记得再加一个参数 map_location=torch.device("cpu")
model = torch.load("network_1.pth")
network = NetWork()
network.load_state_dict(model)
print(network)# 4、将数据输入到模型中进行测试
# 先将数据转换成 pytorch 可以接受的样子，因为 pytorch 中除了图片本身的 C、H、W 之外，还要一个 batch_size
# 因此我们需要将其进行一个转换：batch_size=1, C=3, H=32, W=32
image = torch.reshape(image, (1, 3, 32, 32))
output = network(image)
print(output)

总结

模板基本就是上面这些了，对于不同的学习任务上面的模板可能需要针对性的稍微的进行一点改动，但大体上都是如此，可以多参考 GitHub 上优秀的项目源码。

一定要继续加油嗷！