模型保存、训练与验证

模型保存、训练与验证

网络模型的保存与读取

在 PyTorch 中，模型的保存与加载有两种主要方式：

1. 保存 & 加载完整模型（包括网络结构和参数）
2. 只保存 & 加载模型参数（推荐方式）

可以直接跳过看最后几行。

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 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()
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 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 = 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 x

if __name__ == '__main__':
    # 测试网络的验证正确性
    model = Model()
    input = torch.ones((64, 3, 32, 32))  # batch_size=64, 3通道, 32x32
    output = 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):  # 训练多个 epoch
    print("---------- 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 并记录到 TensorBoard
        if 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 和 batchnorm

    with 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 time

start_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 之间来回交换，也有可能导致变慢。

完整模型验证

测试的核心思路：

1. 加载并预处理测试数据（如图片）

   模型一般要求输入是固定大小的张量，而测试数据可能是各种格式的图片，因此需要进行转换，如 Resize() 、 ToTensor() 。

2. 加载训练好的模型

   需要加载训练好的 pth 模型文件，模型结构必须与训练时一致，否则加载会失败。如果训练时使用了 GPU，推理时在 CPU 运行，则需要 map_location=torch.device('cpu') 进行映射。

3. 进行前向推理

   测试阶段要关闭梯度计算并设置为 eval 模式。还要调整输入维度，添加 batch 维度。

4. 解析并输出预测结果

   多分类问题，需要列出 classes 和 predicted_class ，将数值索引转换成具体的类别名称。

开源项目

开源项目的核心模块

模块	作用
train.py	训练模型的主文件
test.py	测试模型
models/	定义 CycleGAN 和 Pix2Pix 结构
options/	训练/测试的参数
data/	数据加载与预处理
util/	处理可视化、日志

在接触一个新的开源项目时，可以遵循以下步骤：

1. 阅读 README.md
   • 了解项目的用途、安装方法、使用说明。
   • 关注依赖库、数据格式、训练/测试指令等信息。
2. 分析 train.py
   • train.py 是训练脚本，通常包括训练数据加载、模型创建、训练循环、结果可视化、模型保存。
3. 查看 options 目录
   • 该目录下的 train_options.py 定义了训练参数（如学习率、epoch 数、批次大小等）。
   • base_options.py 中定义了通用参数，如数据路径 -dataroot。
4. 分析 models 目录
   • create_model.py 用于加载具体的模型（如 CycleGAN / Pix2Pix）。
   • 具体模型代码通常在 networks.py 或 model_name_model.py。
5. 运行 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 轮。