PyTorch 模型保存与加载
💾 核心概念
PyTorch 中模型的"参数"存储在 state_dict 中------一个 Python 字典，将每一层的名称映射到其参数张量。
python
import torch
import torch.nn as nn
model = nn.Sequential(
nn.Linear(10, 64),
nn.ReLU(),
nn.Linear(64, 1)
)
# 查看 state_dict 的键
print(model.state_dict().keys())
# odict_keys(['0.weight', '0.bias', '2.weight', '2.bias'])
# 查看具体的参数值
print(model.state_dict()['0.weight'].shape) # torch.Size([64, 10])📥 保存与加载模型（推荐方式）
保存
python
# ✅ 推荐：只保存参数（state_dict）
torch.save(model.state_dict(), 'model_weights.pth')
# ⚠️ 不推荐：保存整个模型（依赖代码结构）
torch.save(model, 'entire_model.pth')加载
python
# 1. 先创建相同结构的模型
model = MyModel(input_dim=10, hidden_dim=64, output_dim=1)
# 2. 加载参数
model.load_state_dict(torch.load('model_weights.pth'))
# 3. 切换到评估模式
model.eval()为什么推荐 state_dict？
| 方式 | 优点 | 缺点 |
|---|---|---|
state_dict |
文件小、跨平台、不需源码 | 需知道模型结构 |
| 整个模型 | 加载方便 | 文件大、依赖代码结构、安全性 |
📦 完整训练检查点（Checkpoint）
训练中断时可以恢复，包含模型参数、优化器状态、当前 epoch、损失等。
保存检查点
python
checkpoint = {
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'best_val_acc': best_val_acc,
'loss': loss.item(),
'config': {
'input_dim': 784,
'hidden_dim': 256,
'output_dim': 10,
}
}
torch.save(checkpoint, f'checkpoint_epoch_{epoch}.pth')恢复检查点
python
# 加载
checkpoint = torch.load('checkpoint_epoch_50.pth')
# 恢复模型
model = MyModel(**checkpoint['config'])
model.load_state_dict(checkpoint['model_state_dict'])
# 恢复优化器
optimizer = torch.optim.Adam(model.parameters())
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
# 恢复调度器
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
# 恢复训练状态
start_epoch = checkpoint['epoch'] + 1
best_val_acc = checkpoint['best_val_acc']
# 继续训练
for epoch in range(start_epoch, num_epochs):
train(...)
validate(...)保存最佳模型（训练中）
python
best_acc = 0.0
for epoch in range(num_epochs):
train(...)
val_acc = validate(...)
# 保存最佳
if val_acc > best_acc:
best_acc = val_acc
torch.save(model.state_dict(), 'best_model.pth')
print(f'Saved best model (acc: {best_acc:.4f})')
# 定期保存（每 10 个 epoch）
if (epoch + 1) % 10 == 0:
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'val_acc': val_acc,
}, f'checkpoint_epoch_{epoch+1}.pth')🏷️ 设备映射 --- CPU/GPU 之间
python
# ===== 在 CPU 上加载 GPU 训练的模型 =====
model = MyModel()
model.load_state_dict(torch.load('model.pth', map_location='cpu'))
# ===== 在 GPU 上加载 =====
model = MyModel().to('cuda')
model.load_state_dict(torch.load('model.pth'))
# ===== 加载到指定 GPU =====
model.load_state_dict(torch.load('model.pth', map_location='cuda:1'))
# ===== 自动选择设备 =====
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.load_state_dict(
torch.load('model.pth', map_location=device)
)🎯 迁移学习与微调（Fine-Tuning）
加载预训练模型 + 修改输出层
python
import torchvision.models as models
# 1. 加载预训练模型
model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
# 或旧版写法: models.resnet50(pretrained=True)
# 2. 冻结 backbone（可选）
for param in model.parameters():
param.requires_grad = False
# 3. 替换分类头
num_classes = 10
model.fc = nn.Sequential(
nn.Linear(model.fc.in_features, 512),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(512, num_classes)
)
# 4. 或只冻结部分层
for name, param in model.named_parameters():
if 'layer4' in name or 'fc' in name:
param.requires_grad = True # 只训练 layer4 和 fc
else:
param.requires_grad = False
# 5. 用较小的学习率训练
optimizer = optim.Adam(
filter(lambda p: p.requires_grad, model.parameters()),
lr=1e-4
)加载自己训练的权重做微调
python
# 加载之前训练的权重
pretrained_dict = torch.load('pretrained_weights.pth')
model_dict = model.state_dict()
# 1. 只加载匹配的键（忽略不匹配的）
pretrained_dict = {k: v for k, v in pretrained_dict.items()
if k in model_dict and v.shape == model_dict[k].shape}
# 2. 更新当前模型
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)查看不匹配的层
python
pretrained_dict = torch.load('weights.pth')
model_dict = model.state_dict()
# 忽略的键
ignored = set(pretrained_dict.keys()) - set(model_dict.keys())
# 缺失的键（需要随机初始化的）
missing = set(model_dict.keys()) - set(pretrained_dict.keys())
# 形状不匹配的
mismatched = {k for k in pretrained_dict
if k in model_dict and pretrained_dict[k].shape != model_dict[k].shape}
print(f'Ignored: {ignored}')
print(f'Missing: {missing}')
print(f'Mismatched: {mismatched}')🌐 TorchVision 预训练模型
python
import torchvision.models as models
# ResNet 家族
model = models.resnet18(weights='IMAGENET1K_V1')
model = models.resnet34(weights='DEFAULT')
model = models.resnet50(weights='DEFAULT')
model = models.resnet101(weights='DEFAULT')
# EfficientNet
model = models.efficientnet_b0(weights='DEFAULT')
model = models.efficientnet_v2_s(weights='DEFAULT')
# ViT (Vision Transformer)
model = models.vit_b_16(weights='DEFAULT')
model = models.vit_l_16(weights='DEFAULT')
# MobileNet（移动端）
model = models.mobilenet_v3_small(weights='DEFAULT')
model = models.mobilenet_v3_large(weights='DEFAULT')
# DenseNet
model = models.densenet121(weights='DEFAULT')
# Swin Transformer
model = models.swin_t(weights='DEFAULT')🔄 导出为其他格式
TorchScript（跨语言部署）
python
# 方式 1: 追踪（Tracing）
example_input = torch.randn(1, 3, 224, 224)
traced_model = torch.jit.trace(model, example_input)
traced_model.save('model_traced.pt')
# 加载
model = torch.jit.load('model_traced.pt')
# 方式 2: 脚本（Scripting）
scripted_model = torch.jit.script(model)
scripted_model.save('model_scripted.pt')ONNX（跨框架部署）
python
# 导出 ONNX
torch.onnx.export(
model,
torch.randn(1, 3, 224, 224), # 示例输入
'model.onnx',
input_names=['input'],
output_names=['output'],
dynamic_axes={
'input': {0: 'batch_size'},
'output': {0: 'batch_size'}
},
opset_version=17
)
# 验证 ONNX 模型
import onnx
onnx_model = onnx.load('model.onnx')
onnx.checker.check_model(onnx_model)🧹 清理与压缩
python
# 移除不需要的键（如 optimizer state）
checkpoint = torch.load('full_checkpoint.pth')
model_weights = checkpoint['model_state_dict']
torch.save(model_weights, 'model_only.pth')
# 只保存权重（去除梯度信息）
model.eval()
state_dict = {k: v.detach().cpu() for k, v in model.state_dict().items()}
torch.save(state_dict, 'model_inference.pth')📝 速查表
| 需求 | 代码 |
|---|---|
| 保存参数 | torch.save(model.state_dict(), 'm.pth') |
| 加载参数 | model.load_state_dict(torch.load('m.pth')) |
| 保存检查点 | torch.save({'model': ..., 'opt': ..., 'epoch': e}, 'ckpt.pth') |
| 加载到CPU | torch.load('m.pth', map_location='cpu') |
| 保存最佳模型 | if acc > best: torch.save(...) |
| 加载预训练 | models.resnet50(weights='DEFAULT') |
| 冻结层 | param.requires_grad = False |
| 部分加载权重 | pretrained_dict = {k:v for k,v in ... if k in model_dict} |
| 导出TorchScript | torch.jit.trace(model, input).save('m.pt') |
| 导出ONNX | torch.onnx.export(model, input, 'm.onnx') |
\[pytorch-总览\|← 返回总览\]