模型保存、训练与验证
网络模型的保存与读取
在 PyTorch 中，模型的保存与加载有两种主要方式：
- 保存 & 加载完整模型（包括网络结构和参数）
- 只保存 & 加载模型参数（推荐方式）
可以直接跳过看最后几行。
PyTorch 提供 torch.save() 方法来保存模型，可以选择保存整个模型结构或只保存参数。
保存整个模型
import torch
from torchvision import models# 加载 VGG16 模型（未使用预训练权重）
vgg16 = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1)# 方式1：保存整个模型（包括结构 + 参数）
torch.save(vgg16, "vgg16_method1.pth")
该方法会保存模型的结构和参数。但如果是自定义模型，加载时必须能够访问到模型的定义（否则会报错）。
import torch
from torch import nnclass Model(nn.Module):def __init__(self):super(Model, self).__init__()self.model1 = nn.Sequential(nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1, padding=2),nn.MaxPool2d(kernel_size=2),nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2),nn.MaxPool2d(kernel_size=2),nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding=2),nn.MaxPool2d(kernel_size=2),nn.Flatten(),nn.Linear(in_features=1024, out_features=64),nn.Linear(in_features=64, out_features=10))def forward(self, x):x = self.model1(x)return xmodel = Model()
torch.save(model, "method1.pth")
只保存模型参数
# 方式2（推荐）：仅保存模型参数
torch.save(vgg16.state_dict(), "vgg16_method2.pth")
该方法只保存参数，不包含网络结构。存储文件更小，加载更灵活，但是加载时需要重新定义模型。
加载完整模型
import torch# 加载整个模型（包括结构和参数）
model = torch.load("vgg16_method1.pth")
print(model) # 直接恢复模型
如果是自定义模型，加载时必须在当前代码中定义相同的模型结构，否则会报 AttributeError。
import torch# 加载整个模型（包括结构和参数）
model = torch.load("method1.pth")
print(model) # 报错
上面未重新定义相同的模型结构报错，下面定义后再加载才能成功。
import torch
from torch import nnclass Model(nn.Module):def __init__(self):super(Model, self).__init__()self.model1 = nn.Sequential(nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1, padding=2),nn.MaxPool2d(kernel_size=2),nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2),nn.MaxPool2d(kernel_size=2),nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding=2),nn.MaxPool2d(kernel_size=2),nn.Flatten(),nn.Linear(in_features=1024, out_features=64),nn.Linear(in_features=64, out_features=10))def forward(self, x):x = self.model1(x)return x# 加载整个模型（包括结构和参数）
model = torch.load("method1.pth")
print(model) # 直接恢复模型只加载模型参数
import torch
import torchvision.models as models# 重新构建 VGG16 模型
vgg16 = models.vgg16(weights=None)# 加载模型参数
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))# 现在 `vgg16` 具有之前保存的参数
print(vgg16)该方法不需要存储整个模型结构，节省存储空间。只要模型结构一致，就可以加载到不同的环境中。
import torch
from torch import nn# 重新定义模型（结构必须和原模型一致）
class Model(nn.Module):def __init__(self):super(Model, self).__init__()self.model1 = nn.Sequential(nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1, padding=2),nn.MaxPool2d(kernel_size=2),nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2),nn.MaxPool2d(kernel_size=2),nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding=2),nn.MaxPool2d(kernel_size=2),nn.Flatten(),nn.Linear(in_features=1024, out_features=64),nn.Linear(in_features=64, out_features=10))def forward(self, x):x = self.model1(x)return x# 重新实例化模型
model = Model()
# 加载参数
model.load_state_dict(torch.load("method2.pth"))
print(model)
既然都需要重新加载模型，索性只保存模型参数：
要确保保存和加载时Model类定义必须完全一致，并且必须先实例化再保存或加载模型。
class Model(nn.Module):......
model = Model()
torch.save(model.state_dict(), "model.pth") # 只保存参数class Model(nn.Module):......
model = Model()
model.load_state_dict(torch.load("model.pth")) # 只加载参数
如果是在GPU训练，在CPU加载，需要解决兼容问题。
model.load_state_dict(torch.load("model.pth", **map_location=torch.device('cpu')**))
完整模型训练
以 CIFAR-10 数据集为例，训练一个分类模型。
首先在 model.py 中定义模型
import torch
from torch import nn# 搭建神经网络（10分类网络）
class Model(nn.Module):def __init__(self):super(Model, self).__init__()self.model = nn.Sequential(nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1, padding=2),# 输入是32x32的，输出还是32x32的nn.MaxPool2d(kernel_size=2),nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2),# 输入输出都是16x16的nn.MaxPool2d(kernel_size=2),nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding=2),nn.MaxPool2d(kernel_size=2),nn.Flatten(),nn.Linear(in_features=64 * 4 * 4, out_features=64),nn.Linear(in_features=64, out_features=10))def forward(self, x):x = self.model(x)return xif __name__ == '__main__':# 测试网络的验证正确性model = Model()input = torch.ones((64, 3, 32, 32)) # batch_size=64, 3通道, 32x32output = model(input)print(output.shape)# torch.Size([64, 10])
模型的训练遵循加载数据-定义模型-训练-评估-保存模型的流程。
在 train.py 中训练模型
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torch import nn
from torch.utils.data import DataLoader
from model import * # 从 model.py 文件导入自定义的神经网络 Model# ----------------------- 1. 加载数据集 -----------------------# 下载并加载 CIFAR-10 训练集（自动转换为 tensor）
train_data = torchvision.datasets.CIFAR10(root="Dataset", train=True, transform=torchvision.transforms.ToTensor(),download=True)
# 下载并加载 CIFAR-10 测试集
test_data = torchvision.datasets.CIFAR10(root="Dataset", train=False, transform=torchvision.transforms.ToTensor(),download=True)# 计算训练集和测试集的大小
train_data_size = len(train_data)
test_data_size = len(test_data)# 打印数据集大小
print("Length of train set: {}".format(train_data_size))
print("Length of test set: {}".format(test_data_size))# 使用 DataLoader 进行批量加载数据
train_dataloader = DataLoader(train_data, batch_size=64, shuffle=True) # 训练数据，shuffle=True 让数据打乱
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=False) # 测试数据，顺序读取# ----------------------- 2. 实例化模型、损失函数、优化器 -----------------------# 实例化自定义的神经网络模型
model = Model()# 选择交叉熵损失函数（用于分类任务）
loss_fn = nn.CrossEntropyLoss()# 选择优化器（随机梯度下降 SGD）
learning_rate = 0.01 # 设定学习率
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate) # 传入模型参数# 初始化 TensorBoard 记录器
writer = SummaryWriter("logs")# ----------------------- 3. 训练参数初始化 -----------------------total_train_step = 0 # 记录训练步数
total_test_step = 0 # 记录测试步数
epoch = 10 # 设定训练轮数（整个数据集遍历 10 次）# ----------------------- 4. 开始训练 -----------------------for i in range(epoch): # 训练多个 epochprint("---------- Epoch {} -----------".format(i + 1))# 训练模式model.train()# 遍历训练数据for data in train_dataloader:imgs, targets = data # 获取批量数据outputs = model(imgs) # 前向传播，计算预测值loss = loss_fn(outputs, targets) # 计算损失值# 反向传播优化模型optimizer.zero_grad() # 清空梯度，防止累积loss.backward() # 计算梯度optimizer.step() # 更新参数total_train_step += 1 # 训练步数 +1# 每 100 次迭代打印一次 loss 并记录到 TensorBoardif total_train_step % 100 == 0:print("Train {}, loss: {}".format(total_train_step, loss.item()))writer.add_scalar("train_loss", loss.item(), total_train_step) # 记录训练损失# ----------------------- 5. 测试模型 -----------------------total_test_loss = 0 # 记录测试集上的总损失total_test_accuracy = 0 # 记录测试集上的正确预测数量model.eval() # 切换到测试模式，避免使用 dropout 和 batchnormwith torch.no_grad(): # 关闭梯度计算，加速测试过程for data in test_dataloader:imgs, targets = data # 获取批量测试数据outputs = model(imgs) # 前向传播，得到预测值loss = loss_fn(outputs, targets) # 计算损失total_test_loss += loss.item() # 记录测试总损失accuracy = (outputs.argmax(1) == targets).sum() # 计算正确预测的样本数total_test_accuracy += accuracy # 累加正确预测的样本数total_test_step += 1 # 记录测试轮数# 打印测试结果print("Test set loss: {:.4f}".format(total_test_loss))print("Test set accuracy: {:.2f}%".format((total_test_accuracy / test_data_size) * 100))# 记录测试损失和准确率writer.add_scalar("test_loss", total_test_loss, total_test_step)writer.add_scalar("test_accuracy", total_test_accuracy / test_data_size, total_test_step)# ----------------------- 6. 保存模型 -----------------------torch.save(model.state_dict(), "model_{}.pth".format(i)) # 每轮训练后保存模型参数print("Model saved.")# ----------------------- 7. 关闭 TensorBoard 记录器 -----------------------
writer.close()GPU训练
PyTorch 主要有两种方式让模型在 GPU 上运行：
方法 1：直接将网络模型、数据和损失函数转移到 GPU
在创建模型、数据和损失函数部分，将 cuda() 添加到相应的对象上：
# 1. 加载模型并移动到 GPU
model = Model()
model = model.cuda() # 直接将模型移动到 GPU# 2. 加载数据并移动到 GPU
inputs, labels = data
inputs = inputs.cuda() # 输入数据移动到 GPU
labels = labels.cuda() # 目标数据移动到 GPU# 3. 定义损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.cuda()
但是这种方法在 CPU 和 GPU 之间切换不够灵活，需要手动修改代码。
方法 2：使用 torch.device() 控制 GPU 训练
PyTorch 提供 torch.device()，可以动态选择运行设备，代码在 CPU 和 GPU 上都能运行：
# 1. 检查 GPU 是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 2. 将模型移动到 GPU（或者 CPU）
model = Model().to(device)# 3. 将数据移动到 GPU
inputs, labels = inputs.to(device), labels.to(device)# 4. 将损失函数移动到 GPU
loss_fn = nn.CrossEntropyLoss().to(device)# 4. 训练过程
outputs = model(inputs) # 计算仍然在 GPU 上
loss = loss_fn (outputs, labels) # 计算损失
loss.backward()
optimizer.step()为了测试 GPU 是否加速，可以使用 time 模块来计时：
import timestart_time = time.time()
# ......
end_time = time.time()print("Training Time: {:.2f} seconds".format(end_time - start_time))
但有时 GPU 训练比 CPU 慢，可能是因为 CPU 到 GPU 的数据传输存在开销，可以提前把数据全部加载到 GPU 解决。如果模型很小，计算量不足以让 GPU 发挥优势，反而可能拖慢速度。如果 GPU 显存不够，数据可能在 CPU 和 GPU 之间来回交换，也有可能导致变慢。
完整模型验证
测试的核心思路：
-
加载并预处理测试数据（如图片）
模型一般要求输入是固定大小的张量，而测试数据可能是各种格式的图片，因此需要进行转换，如
Resize()、ToTensor()。 -
加载训练好的模型
需要加载训练好的
pth模型文件，模型结构必须与训练时一致，否则加载会失败。如果训练时使用了 GPU，推理时在 CPU 运行，则需要map_location=torch.device('cpu')进行映射。 -
进行前向推理
测试阶段要关闭梯度计算并设置为
eval模式。还要调整输入维度，添加 batch 维度。 -
解析并输出预测结果
多分类问题，需要列出
classes和predicted_class，将数值索引转换成具体的类别名称。
开源项目
开源项目的核心模块
| 模块 | 作用 |
|---|---|
train.py | 训练模型的主文件 |
test.py | 测试模型 |
models/ | 定义 CycleGAN 和 Pix2Pix 结构 |
options/ | 训练/测试的参数 |
data/ | 数据加载与预处理 |
util/ | 处理可视化、日志 |
在接触一个新的开源项目时，可以遵循以下步骤：
- 阅读
README.md- 了解项目的用途、安装方法、使用说明。
- 关注依赖库、数据格式、训练/测试指令等信息。
- 分析
train.pytrain.py是训练脚本，通常包括训练数据加载、模型创建、训练循环、结果可视化、模型保存。
- 查看
options目录- 该目录下的
train_options.py定义了训练参数（如学习率、epoch 数、批次大小等）。 base_options.py中定义了通用参数，如数据路径-dataroot。
- 该目录下的
- 分析
models目录create_model.py用于加载具体的模型（如 CycleGAN / Pix2Pix）。- 具体模型代码通常在
networks.py或model_name_model.py。
- 运行
test.py- 该脚本通常用于推理/验证，可以输入一张图片，输出转换后的图像。
训练参数在 train_options.py 文件中定义：
parser.add_argument('--display_freq', type=int, default=400, help='显示训练结果的频率')
parser.add_argument('--save_latest_freq', type=int, default=5000, help='保存最新模型的频率')
parser.add_argument('--save_epoch_freq', type=int, default=5, help='每多少个 epoch 保存一次')
parser.add_argument('--n_epochs', type=int, default=100, help='初始学习率训练多少个 epoch')
parser.add_argument('--n_epochs_decay', type=int, default=100, help='学习率衰减多少个 epoch')
parser.add_argument('--lr', type=float, default=0.0002, help='学习率')
parser.add_argument('--gan_mode', type=str, default='lsgan', help='GAN 损失模式，可选 vanilla | lsgan | wgangp')这些参数用于 train.py，可以在运行时通过命令行修改：
python train.py --dataroot ./datasets/maps --name maps_cyclegan --model cycle_gan --n_epochs 200
-dataroot ./datasets/maps：指定数据集路径。-name maps_cyclegan：实验名称，用于日志保存。-model cycle_gan：使用 CycleGAN 模型。-n_epochs 200：训练 200 轮。