PyTorch模型保存方式

为啥全要学2025-05-20 13:15

PyTorch提供两种主流模型保存方式和一种训练断点保存与恢复的方法。

1. 仅保存模型参数（推荐）

python 复制代码

# 保存
torch.save(model.state_dict(), "model_params.pth")  

# 加载
new_model = TheModelClass()  
new_model.load_state_dict(torch.load("model_params.pth"))
new_model.eval()

核心优势：

文件体积小（仅参数数据）
避免PyTorch版本兼容问题
支持跨模型结构迁移（需设置strict=False）

2. 保存完整模型对象

python 复制代码

# 保存
torch.save(model, "full_model.pth")  

# 加载   loaded_model = torch.load("full_model.pth")
loaded_model.eval()

适用场景：

快速原型验证
模型结构包含动态逻辑（如自定义前向传播）

3. 训练断点保存与恢复

python 复制代码

# 保存检查点
checkpoint = {
    'epoch': current_epoch,
    'model_state': model.state_dict(),
    'optimizer_state': optimizer.state_dict(),
    'loss': loss_value
}
torch.save(checkpoint, "checkpoint.tar")

# 恢复训练
model = TheModelClass()
optimizer = torch.optim.Adam(model.parameters())
checkpoint = torch.load("checkpoint.tar")
model.load_state_dict(checkpoint['model_state'])
optimizer.load_state_dict(checkpoint['optimizer_state'])
model.train()  # 保持训练模式

关键细节：

推荐使用.tar后缀区分普通参数文件
自动恢复学习率调度器等训练状态