模型保存与加载

模型的保存与加载是深度学习工程中至关重要的一环。无论是中断恢复、模型部署还是迁移学习，都离不开对模型持久化的深入理解。本文将全面介绍 PyTorch 中 state_dict、Checkpoint、TorchScript、ONNX 等核心概念与实践。

state_dict 概念

什么是 state_dict

在 PyTorch 中，torch.nn.Module 模型的学习参数（权重和偏置）通过 model.parameters() 获取，而模型的全部可学习状态存储在一个有序字典（OrderedDict）中，这就是 state_dict。

python

import torch
import torch.nn as nn

# 定义一个简单模型
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

model = SimpleNet()
print(type(model.state_dict()))  # <class 'collections.OrderedDict'>

# 查看 state_dict 中的键
for key, value in model.state_dict().items():
    print(f"{key}: {value.shape}")
# 输出：
# fc1.weight: torch.Size([256, 784])
# fc1.bias: torch.Size([256])
# fc2.weight: torch.Size([10, 256])
# fc2.bias: torch.Size([10])

为什么推荐使用 state_dict

使用 state_dict 保存和加载模型是 PyTorch 的最佳实践，原因如下：

优势	说明
灵活性	可以只保存参数，不依赖模型类的定义代码
兼容性	模型结构变更后，仍可部分加载匹配的参数
安全性	不使用 Python pickle 反序列化整个对象，更安全
轻量	只包含参数数据，不包含优化器状态等额外信息
跨设备	方便在 CPU/GPU 之间切换
社区标准	Hugging Face、TIMM 等主流库均使用此方式

保存与加载模型权重

torch.save 与 torch.load

python

# ==================== 保存模型权重 ====================
# 推荐方式：只保存 state_dict
torch.save(model.state_dict(), 'model_weights.pth')

# 也可以保存为 .pt 格式（本质相同）
torch.save(model.state_dict(), 'model_weights.pt')

# ==================== 加载模型权重 ====================
# 1. 先实例化模型结构
model = SimpleNet()

# 2. 加载 state_dict
model.load_state_dict(torch.load('model_weights.pth'))

# 3. 设置为评估模式（重要！）
model.eval()

strict 参数

当模型结构与保存的权重不完全匹配时（如微调时替换了最后一层），可以使用 strict=False：

python

# 假设我们修改了模型的最后分类层
class FineTunedNet(nn.Module):
    def __init__(self, num_classes=5):  # 原来是 10 类，现在改为 5 类
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(256, num_classes)  # 不同的输出维度

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

model = FineTunedNet(num_classes=5)

# strict=False: 忽略不匹配的键（如 fc2 的维度不同）
missing, unexpected = model.load_state_dict(
    torch.load('model_weights.pth'), strict=False
)
print(f"Missing keys: {missing}")      # ['fc2.weight', 'fc2.bias']
print(f"Unexpected keys: unexpected")  # 可能为空或列出多余键

# fc1 的权重正常加载，fc2 保持随机初始化

保存完整模型 vs 保存 state_dict

两种方式对比

python

# ==================== 方式一：保存完整模型（不推荐） ====================
torch.save(model, 'full_model.pth')
loaded_model = torch.load('full_model.pth')

# ==================== 方式二：保存 state_dict（推荐） ====================
torch.save(model.state_dict(), 'model_weights.pth')
model = SimpleNet()  # 必须先创建模型实例
model.load_state_dict(torch.load('model_weights.pth'))

对比项	保存完整模型	保存 state_dict
序列化方式	Python pickle 序列化整个对象	只保存参数字典
文件依赖	依赖定义模型的源代码文件	只需重新定义模型类
灵活性	低，无法部分加载	高，支持 partial loading
安全性	较低（pickle 可执行任意代码）	较高
文件大小	较大（包含整个对象图）	较小（只有参数）
推荐程度	不推荐	强烈推荐

Checkpoint 机制

Checkpoint（检查点）不仅保存模型权重，还保存训练状态（epoch、优化器、调度器等），用于实现断点续训。

完整的 Checkpoint 保存与加载

python

import torch
import torch.nn as nn
import torch.optim as optim

# ==================== 训练过程中保存 Checkpoint ====================
def save_checkpoint(model, optimizer, scheduler, epoch, best_acc, train_losses, path):
    """保存训练检查点"""
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict() if scheduler else None,
        'best_acc': best_acc,
        'train_losses': train_losses,
        # 可以添加任何需要恢复的信息
        'torch_rng_state': torch.get_rng_state(),
        'cuda_rng_state': torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
    }
    torch.save(checkpoint, path)
    print(f"Checkpoint saved to {path} (epoch {epoch}, best_acc {best_acc:.2f}%)")

# ==================== 从 Checkpoint 恢复训练 ====================
def load_checkpoint(path, model, optimizer=None, scheduler=None):
    """加载训练检查点"""
    checkpoint = torch.load(path, weights_only=False)

    model.load_state_dict(checkpoint['model_state_dict'])

    if optimizer and 'optimizer_state_dict' in checkpoint:
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

    if scheduler and checkpoint.get('scheduler_state_dict'):
        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])

    # 恢复随机状态（确保可复现性）
    if 'torch_rng_state' in checkpoint:
        torch.set_rng_state(checkpoint['torch_rng_state'])
    if checkpoint.get('cuda_rng_state') and torch.cuda.is_available():
        torch.cuda.set_rng_state_all(checkpoint['cuda_rng_state'])

    epoch = checkpoint['epoch']
    best_acc = checkpoint.get('best_acc', 0.0)
    train_losses = checkpoint.get('train_losses', [])

    print(f"Checkpoint loaded from {path} (epoch {epoch}, best_acc {best_acc:.2f}%)")
    return epoch, best_acc, train_losses

完整的断点续训流程

python

# ==================== 完整的断点续训示例 ====================
import os

# 初始化
model = SimpleNet()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
criterion = nn.CrossEntropyLoss()

# 训练配置
num_epochs = 100
start_epoch = 0
best_acc = 0.0
train_losses = []
checkpoint_path = 'checkpoint.pth'
best_model_path = 'best_model.pth'

# 尝试加载检查点（断点续训）
if os.path.exists(checkpoint_path):
    start_epoch, best_acc, train_losses = load_checkpoint(
        checkpoint_path, model, optimizer, scheduler
    )
    start_epoch += 1  # 从下一个 epoch 开始
    print(f"Resuming training from epoch {start_epoch}")

# 训练循环
for epoch in range(start_epoch, num_epochs):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    for inputs, targets in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

    scheduler.step()
    train_acc = 100.0 * correct / total
    avg_loss = running_loss / len(train_loader)
    train_losses.append(avg_loss)

    print(f"Epoch [{epoch}/{num_epochs}] Loss: {avg_loss:.4f} Acc: {train_acc:.2f}%")

    # 验证
    val_acc = evaluate(model, val_loader)

    # 保存最优模型
    if val_acc > best_acc:
        best_acc = val_acc
        torch.save(model.state_dict(), best_model_path)
        print(f"  New best model saved! Accuracy: {best_acc:.2f}%")

    # 每 N 个 epoch 或每个 epoch 保存 checkpoint
    save_checkpoint(model, optimizer, scheduler, epoch, best_acc, train_losses, checkpoint_path)

print(f"Training finished. Best accuracy: {best_acc:.2f}%")

保存多个最佳模型（Top-K）

python

import heapq

class TopKCheckpointManager:
    """保存验证指标最优的 K 个模型"""

    def __init__(self, save_dir='checkpoints', k=3, mode='max'):
        self.save_dir = save_dir
        self.k = k
        self.mode = mode
        self.checkpoints = []  # (score, filename) 的列表
        os.makedirs(save_dir, exist_ok=True)

    def update(self, score, model, epoch, **extra_info):
        filename = f"checkpoint_epoch{epoch}_score{score:.4f}.pth"
        filepath = os.path.join(self.save_dir, filename)

        # 保存模型
        state = {
            'epoch': epoch,
            'score': score,
            'model_state_dict': model.state_dict(),
            **extra_info
        }
        torch.save(state, filepath)

        # 管理 top-k
        entry = (score if self.mode == 'max' else -score, filename)
        self.checkpoints.append(entry)

        if len(self.checkpoints) > self.k:
            # 移除最差的 checkpoint
            self.checkpoints.sort()
            _, worst_file = self.checkpoints.pop(0)
            worst_path = os.path.join(self.save_dir, worst_file)
            if os.path.exists(worst_path):
                os.remove(worst_path)
                print(f"  Removed checkpoint: {worst_file}")

        print(f"  Saved checkpoint: {filename} (score: {score:.4f})")

# 使用
topk_manager = TopKCheckpointManager(save_dir='checkpoints', k=3, mode='max')
# 在验证循环中
topk_manager.update(val_acc, model, epoch, optimizer_state=optimizer.state_dict())

跨设备加载

GPU → CPU

python

# 场景：模型在 GPU 上训练，需要在 CPU 上加载

# 方式一：指定 map_location
model.load_state_dict(torch.load('model_gpu.pth', map_location='cpu'))

# 方式二：加载后再转设备
model.load_state_dict(torch.load('model_gpu.pth', map_location=torch.device('cpu')))

# 对于完整 checkpoint
checkpoint = torch.load('checkpoint.pth', map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])

CPU → GPU

python

# 场景：模型在 CPU 上保存，需要加载到 GPU

# 方式一：直接加载到指定 GPU
model.load_state_dict(torch.load('model_cpu.pth', map_location='cuda:0'))
model = model.to('cuda:0')

# 方式二：先加载到 CPU，再转移到 GPU
model.load_state_dict(torch.load('model_cpu.pth', map_location='cpu'))
model = model.cuda()

多 GPU（DataParallel）→ 单 GPU

使用 DataParallel 训练的模型，其 state_dict 的键名会有 module. 前缀。加载到单 GPU 模型时需要处理：

python

# ==================== 多 GPU 保存的模型加载到单 GPU ====================
# DataParallel 保存的 state_dict 键名格式：
# module.fc1.weight, module.fc1.bias, module.fc2.weight, ...

# 方式一：去掉 module. 前缀
def remove_module_prefix(state_dict):
    """去掉 DataParallel 添加的 module. 前缀"""
    new_state_dict = {}
    for key, value in state_dict.items():
        new_key = key.replace('module.', '') if key.startswith('module.') else key
        new_state_dict[new_key] = value
    return new_state_dict

checkpoint = torch.load('dp_model.pth', map_location='cpu')
if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
    state_dict = checkpoint['model_state_dict']
else:
    state_dict = checkpoint

state_dict = remove_module_prefix(state_dict)
model.load_state_dict(state_dict)

# 方式二：使用 strict=False 并手动处理
# 如果模型结构完全对应，可以用更简洁的方式
from collections import OrderedDict

state_dict = torch.load('dp_model.pth', map_location='cpu')
new_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k[7:] if k.startswith('module.') else k  # 去掉 'module.'
    new_state_dict[name] = v

model.load_state_dict(new_state_dict)

# ==================== 单 GPU 加载到 DataParallel ====================
model = nn.DataParallel(model)
# 如果原始权重没有 module. 前缀，DataParallel 会自动添加
model.load_state_dict(torch.load('single_gpu_model.pth'))

# ==================== DDP（DistributedDataParallel）模型处理 ====================
# DDP 也会添加 module. 前缀，处理方式与 DataParallel 相同
def load_ddp_checkpoint_to_single_model(checkpoint_path, model):
    checkpoint = torch.load(checkpoint_path, map_location='cpu')
    state_dict = checkpoint['model_state_dict']

    # 去掉 module. 前缀
    from collections import OrderedDict
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k.replace('module.', '')
        new_state_dict[name] = v

    model.load_state_dict(new_state_dict)
    return model

自动跨设备加载（通用方案）

python

def load_model_flexible(path, model, device='cpu'):
    """通用的模型加载函数，自动处理跨设备和多 GPU 键名"""
    checkpoint = torch.load(path, map_location=device, weights_only=False)

    # 支持直接的 state_dict 或包含 state_dict 的 checkpoint
    if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
        state_dict = checkpoint['model_state_dict']
    elif isinstance(checkpoint, dict) and any(isinstance(v, torch.Tensor) for v in checkpoint.values()):
        state_dict = checkpoint
    else:
        raise ValueError("Unsupported checkpoint format")

    # 去掉 module. 前缀（如果存在）
    from collections import OrderedDict
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        new_key = k.replace('module.', '')
        new_state_dict[new_key] = v

    # 加载（允许部分匹配）
    missing, unexpected = model.load_state_dict(new_state_dict, strict=False)
    if missing:
        print(f"Warning: Missing keys: {missing}")
    if unexpected:
        print(f"Warning: Unexpected keys: {unexpected}")

    model = model.to(device)
    return model, checkpoint

模型导出

训练好的模型通常需要导出为特定格式以满足部署需求。PyTorch 提供了多种导出方式。

TorchScript

TorchScript 是 PyTorch 模型的中间表示（IR），可以脱离 Python 运行环境独立执行。

torch.jit.trace

trace 通过运行一次模型并记录所有操作来生成 TorchScript 模型，适用于不包含数据依赖控制流的模型。

python

import torch

model = SimpleNet()
model.eval()

# 创建示例输入
example_input = torch.randn(1, 784)

# 使用 trace 导出
traced_model = torch.jit.trace(model, example_input)

# 保存
traced_model.save('model_traced.pt')

# 加载和推理
loaded_model = torch.jit.load('model_traced.pt')
output = loaded_model(example_input)

trace 的局限性：

python

# 以下情况 trace 无法正确捕获：
class ModelWithControlFlow(nn.Module):
    def forward(self, x):
        if x.sum() > 0:        # 数据依赖的条件分支 —— trace 无法捕获
            return x * 2
        else:
            return x * 3

# trace 只会记录一条执行路径
traced = torch.jit.trace(ModelWithControlFlow(), torch.tensor([1.0]))
print(traced(torch.tensor([-1.0])))  # 输出 tensor([−2.])，错误！应该是 tensor([−3.])

torch.jit.script

script 通过解析 Python 源代码来生成 TorchScript 模型，可以正确处理控制流。

python

class ModelWithControlFlow(nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x * 2
        else:
            return x * 3

model = ModelWithControlFlow()
scripted_model = torch.jit.script(model)
scripted_model.save('model_scripted.pt')

# 验证正确性
print(scripted_model(torch.tensor([-1.0])))  # tensor([-3.])，正确！

trace vs script 对比

特性	`torch.jit.trace`	`torch.jit.script`
原理	运行一次，记录操作	解析 Python 源码
控制流	不支持数据依赖的控制流	完整支持
print/断言	不捕获	可捕获
Python 函数	可以 trace 调用的函数	必须是 TorchScript 兼容的
速度	trace 本身很快	解析可能较慢
适用场景	CNN 等固定计算图模型	RNN、Transformer 等含控制流模型

trace + script 组合使用

python

class Encoder(nn.Module):
    """简单的编码器，不含控制流"""
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 64, 3, padding=1)

    def forward(self, x):
        return torch.relu(self.conv(x))

class Decoder(nn.Module):
    """解码器，包含控制流"""
    def forward(self, x):
        if x.dim() == 4:
            x = x.mean(dim=[2, 3])
        return x

class CombinedModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()

    @torch.jit.export
    def encode(self, x):
        return self.encoder(x)

    def forward(self, x):
        features = self.encoder(x)
        return self.decoder(features)

# 对于混合模型，可以先 script 整体
model = CombinedModel()
scripted = torch.jit.script(model)
scripted.save('combined_model.pt')

ONNX 导出

ONNX（Open Neural Network Exchange）是跨框架的模型交换格式，可以在多种推理引擎上运行。

基本导出

python

import torch.onnx

model = SimpleNet()
model.eval()

# 创建虚拟输入（dummy input）
dummy_input = torch.randn(1, 784)

# 导出为 ONNX
torch.onnx.export(
    model,                           # 要导出的模型
    dummy_input,                     # 虚拟输入
    "model.onnx",                    # 输出文件路径
    export_params=True,              # 导出模型参数
    opset_version=17,                # ONNX opset 版本（推荐 17+）
    do_constant_folding=True,        # 是否执行常量折叠优化
    input_names=['input'],           # 输入名称
    output_names=['output'],         # 输出名称
    dynamic_axes={                   # 动态维度
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'},
    }
)

print("Model exported to model.onnx")

验证导出的 ONNX 模型

python

import onnx
import onnxruntime as ort
import numpy as np

# 1. 验证 ONNX 模型结构
onnx_model = onnx.load('model.onnx')
onnx.checker.check_model(onnx_model)
print("ONNX model is valid!")

# 2. 使用 onnxruntime 推理
session = ort.InferenceSession('model.onnx')

# 准备输入数据
input_data = np.random.randn(1, 784).astype(np.float32)

# 推理
outputs = session.run(None, {'input': input_data})
print(f"ONNX output shape: {outputs[0].shape}")

# 3. 对比 PyTorch 和 ONNX 的输出
model.eval()
with torch.no_grad():
    pytorch_output = model(torch.from_numpy(input_data)).numpy()

diff = np.abs(pytorch_output - outputs[0]).max()
print(f"Max absolute difference: {diff:.8f}")  # 应该非常小（< 1e-5）

导出卷积网络到 ONNX

python

import torchvision

# 使用预训练的 ResNet-18
model = torchvision.models.resnet18(pretrained=True)
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    opset_version=17,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'},
    }
)

# 使用 ONNX Runtime 推理
session = ort.InferenceSession("resnet18.onnx")
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
outputs = session.run(None, {'input': input_data})
print(f"ResNet-18 ONNX output: {outputs[0].shape}")

Torch.export（PyTorch 2.x 新方式）

torch.export 是 PyTorch 2.x 引入的新一代模型导出 API，提供更完整的图捕获能力。

python

# 需要 PyTorch 2.1+
import torch

model = SimpleNet()
model.eval()

# 使用 torch.export 导出
example_input = (torch.randn(1, 784),)
exported_model = torch.export.export(model, example_input)

# 导出的模型是 ExportedProgram 类型
print(type(exported_model))  # <class 'torch.export.exported_program.ExportedProgram'>

# 可以像普通模型一样调用
output = exported_model(torch.randn(2, 784))

# 保存和加载
torch.export.save(exported_model, 'exported_model.pt2')
loaded = torch.export.load('exported_model.pt2')
output = loaded(torch.randn(1, 784))

torch.export 与 torch.jit 的对比

特性	`torch.export`	`torch.jit.trace`	`torch.jit.script`
图捕获	完整的 FX 图	只记录执行路径	解析 Python 源码
控制流	支持（通过特化）	不支持	支持
数据依赖形状	支持	不支持	部分支持
Python 版本兼容	好	好	对语法有严格要求
推荐程度（PyTorch 2.x）	推荐	仅限简单模型	需要控制流时
维护状态	积极开发中	稳定，不再新功能	稳定

高级导出选项

python

import torch

class DynamicModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 5)

    def forward(self, x, flag: bool):
        x = self.linear(x)
        if flag:
            return x * 2
        return x

model = DynamicModel()
example_input = (torch.randn(1, 10), True)

# torch.export 可以处理带控制流的模型
exported = torch.export.export(model, example_input)
print(exported)

# 查看导出的图
print(exported.graph)

安全注意事项

torch.load 的 weights_only 参数

从 PyTorch 1.13 开始，torch.load 引入了 weights_only 参数来增强安全性。

python

# ==================== 安全风险 ====================
# 默认情况下，torch.load 使用 pickle 反序列化，可能执行恶意代码
# 如果加载了不受信任的文件，攻击者可以执行任意代码

# ==================== 安全做法 ====================
# 方式一：weights_only=True（推荐）
# 只允许加载张量、基本类型和 OrderedDict，拒绝执行任意代码
model.load_state_dict(torch.load('model.pth', weights_only=True))

# 方式二：验证文件来源
# 只加载来自可信来源的模型文件
import hashlib

def verify_checksum(filepath, expected_hash):
    """验证文件完整性"""
    with open(filepath, 'rb') as f:
        file_hash = hashlib.sha256(f.read()).hexdigest()
    return file_hash == expected_hash

if verify_checksum('model.pth', 'abc123...'):
    model.load_state_dict(torch.load('model.pth', weights_only=True))
else:
    raise ValueError("File checksum mismatch!")

# ==================== PyTorch 2.6+ 的默认行为 ====================
# 从 PyTorch 2.6 开始，torch.load 默认使用 weights_only=True
# 如果需要加载包含非张量数据的 checkpoint，必须显式设置 weights_only=False
checkpoint = torch.load('checkpoint.pth', weights_only=False)

weights_only 的限制

python

# weights_only=True 时，以下类型可以安全加载：
# - torch.Tensor
# - int, float, str, bool, None
# - list, tuple, dict, OrderedDict
# - set, frozenset

# 以下类型会被拒绝（会抛出异常）：
# - 自定义类实例
# - numpy 数组（需要额外处理）
# - 函数、lambda 表达式

# 如果你的 checkpoint 包含 numpy 数组，需要自定义安全全局变量
import numpy as np

# 方式一：使用 torch.serialization.add_safe_globals
torch.serialization.add_safe_globals([np.core.multiarray._reconstruct])

# 方式二：使用 safe_globals 上下文管理器
from torch.serialization import safe_globals
with safe_globals([np.core.multiarray._reconstruct]):
    checkpoint = torch.load('checkpoint.pth', weights_only=True)

常见错误与解决方案

错误一：KeyError 或 Missing Keys

python

# 错误信息：
# RuntimeError: Error(s) in loading state_dict for SimpleNet:
#   Missing key(s) in state_dict: "fc3.weight", "fc3.bias"
#   Unexpected key(s) in state_dict: "fc2.weight", "fc2.bias"

# 原因：模型结构与保存的权重不匹配（如添加了新层）

# 解决方案：
# 1. 使用 strict=False
missing, unexpected = model.load_state_dict(
    torch.load('model.pth', weights_only=True), strict=False
)
print(f"Missing: {missing}")
print(f"Unexpected: {unexpected}")

# 2. 确保模型定义与保存时完全一致
# 检查模型结构
for name, param in model.named_parameters():
    print(f"{name}: {param.shape}")

错误二：设备不匹配

python

# 错误信息：
# RuntimeError: Expected all tensors to be on the same device,
# but found at least two devices, cuda:0 and cpu!

# 解决方案：使用 map_location
model.load_state_dict(
    torch.load('model.pth', map_location='cpu', weights_only=True)
)

# 或者在加载后转移到目标设备
state_dict = torch.load('model.pth', map_location='cpu', weights_only=True)
model.load_state_dict(state_dict)
model = model.to('cuda:0')

错误三：DataParallel 键名不匹配

python

# 错误信息：
# RuntimeError: Error(s) in loading state_dict for SimpleNet:
#   Unexpected key(s) in state_dict: "module.fc1.weight", "module.fc1.bias"

# 解决方案：去掉 module. 前缀
state_dict = torch.load('model.pth', map_location='cpu', weights_only=True)
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
    new_state_dict[k.replace('module.', '')] = v
model.load_state_dict(new_state_dict)

错误四：版本兼容性问题

python

# 错误信息：
# RuntimeError: version_ <= kMaxSupportedFileFormatVersion ASSERT FAILED

# 原因：用较新版本 PyTorch 保存的模型无法用旧版本加载

# 解决方案：
# 1. 升级 PyTorch 版本
# 2. 用 TorchScript 或 ONNX 格式导出（跨版本兼容性更好）
# 3. 保存时指定 pickle 协议版本（有限帮助）
torch.save(model.state_dict(), 'model.pth', _use_new_zipfile_serialization=True)

错误五：ONNX 导出失败

python

# 错误信息：
# RuntimeError: Exporting the operator 'aten::xxx' to ONNX opset version XX
# is not supported.

# 解决方案：
# 1. 升级 opset_version
torch.onnx.export(model, dummy_input, "model.onnx", opset_version=17)

# 2. 重写不支持的操作
class MyModel(nn.Module):
    def forward(self, x):
        # 不支持的操作可以用等价操作替代
        # 例如：torch.where → torch.where 的 ONNX 支持版本
        return torch.clamp(x, min=0)  # 用 clamp 替代某些不支持的激活函数

# 3. 使用自定义 ONNX 算子（高级用法）
# 参考 torch.onnx.register_custom_op_library

错误六：Checkpoint 加载后 loss 突然增大

python

# 原因：忘记恢复优化器或调度器的状态

# 正确做法：确保恢复所有相关状态
checkpoint = torch.load('checkpoint.pth', weights_only=False)

model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])  # 必须恢复！
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])  # 必须恢复！
start_epoch = checkpoint['epoch'] + 1

# 如果不恢复 optimizer 状态，学习率和动量会重置
# 如果不恢复 scheduler 状态，学习率调度会从头开始

错误七：torch.jit.trace 丢失动态行为

python

# 问题：traced 模型的行为与原始模型不一致

# 诊断方法
model = MyModelWithIf()
example = torch.randn(1, 10)

# 检查 trace 的警告
with torch.jit.optimized_execution(True):
    traced = torch.jit.trace(model, example)
    # 如果有警告："Converting a tensor to a Python boolean might cause the trace
    # to be incorrect"，说明 trace 无法捕获数据依赖的控制流

# 解决方案：改用 torch.jit.script
scripted = torch.jit.script(model)

数据结构

数据库

离散数学

数字电路

概率论

数学建模

软件工程

面向对象程序设计

Python 与数据分析

计算机组成

操作系统

计算机网络

编译原理

图形学

自然语言处理

Linux 实验

分布式系统

Internet原理与技术

BlogWeb

Stars Saying

黑马商城

苍穹外卖

黑马点评

黑马商城单体版

JavaWeb AI Tlias

RuoYi DKD

LinkQwen

InterX

LeetCode Agent Lab

00-预备知识

01-PR1108-逐代码段解析

02-PR1135-Cube化逐代码段解析

03-两个Bug深度剖析

04-开发心得与数值性能实测

01-PR1169-模型推理性能优化

02-PR1194-固定Shape及分档推理适配

NumPy

PyTorch

Python 八股文

LeetCUDA

LeetGPU

Triton

模型保存与加载 ​

state_dict 概念 ​

什么是 state_dict ​

为什么推荐使用 state_dict ​

保存与加载模型权重 ​

torch.save 与 torch.load ​

strict 参数 ​

保存完整模型 vs 保存 state_dict ​

两种方式对比 ​

Checkpoint 机制 ​

完整的 Checkpoint 保存与加载 ​

完整的断点续训流程 ​

保存多个最佳模型（Top-K） ​

跨设备加载 ​

GPU → CPU ​

CPU → GPU ​

多 GPU（DataParallel）→ 单 GPU ​

自动跨设备加载（通用方案） ​

模型导出 ​

TorchScript ​

torch.jit.trace ​

torch.jit.script ​

trace vs script 对比 ​

trace + script 组合使用 ​

ONNX 导出 ​

基本导出 ​

验证导出的 ONNX 模型 ​

导出卷积网络到 ONNX ​

Torch.export（PyTorch 2.x 新方式） ​

torch.export 与 torch.jit 的对比 ​

高级导出选项 ​

安全注意事项 ​

torch.load 的 weights_only 参数 ​

weights_only 的限制 ​

常见错误与解决方案 ​

错误一：KeyError 或 Missing Keys ​

错误二：设备不匹配 ​

错误三：DataParallel 键名不匹配 ​

模型保存与加载

state_dict 概念

什么是 state_dict

为什么推荐使用 state_dict

保存与加载模型权重

torch.save 与 torch.load

strict 参数

保存完整模型 vs 保存 state_dict

两种方式对比

Checkpoint 机制

完整的 Checkpoint 保存与加载

完整的断点续训流程

保存多个最佳模型（Top-K）

跨设备加载

GPU → CPU

CPU → GPU

多 GPU（DataParallel）→ 单 GPU

自动跨设备加载（通用方案）

模型导出

TorchScript

torch.jit.trace

torch.jit.script

trace vs script 对比

trace + script 组合使用

ONNX 导出

基本导出

验证导出的 ONNX 模型

导出卷积网络到 ONNX

Torch.export（PyTorch 2.x 新方式）

torch.export 与 torch.jit 的对比

高级导出选项

安全注意事项

torch.load 的 weights_only 参数

weights_only 的限制

常见错误与解决方案

错误一：KeyError 或 Missing Keys

错误二：设备不匹配

错误三：DataParallel 键名不匹配

错误四：版本兼容性问题

错误五：ONNX 导出失败

错误六：Checkpoint 加载后 loss 突然增大

错误七：torch.jit.trace 丢失动态行为

参考资料