CentOS上PyTorch模型保存与加载方法

在机器学习和深度学习项目中，PyTorch因其灵活性和易用性成为热门框架之一。当我们在CentOS服务器上训练好模型后，如何高效地保存和加载模型成为关键问题。在CentOS系统上使用PyTorch保存和加载模型的完整流程，涵盖常见方法、注意事项以及实际代码示例，帮助开发者快速掌握这一核心技能。

1. PyTorch模型保存的基本方法

PyTorch提供了两种主要的模型保存方式：

1.1 保存整个模型

使用torch.save()直接保存整个模型对象（包括结构和参数）：

import torch
model = ...  # 训练好的模型
torch.save(model, 'model.pth')

优点：加载时无需重新定义模型结构，代码简洁。
缺点：文件较大，且对代码环境依赖性强（需保证模型类定义一致）。

1.2 仅保存模型参数

推荐只保存模型的状态字典（state_dict），更轻量且灵活：

torch.save(model.state_dict(), 'model_weights.pth')

加载时需先实例化模型结构，再加载参数：

model = MyModel()  # 需提前定义模型类
model.load_state_dict(torch.load('model_weights.pth'))

2. 模型加载的注意事项

2.1 设备兼容性

若模型在GPU训练但需在CPU加载，需指定map_location：

model.load_state_dict(torch.load('model_weights.pth', map_location=torch.device('cpu')))

2.2 版本一致性

PyTorch版本差异可能导致加载失败，建议通过以下命令检查环境：

pip list | grep torch  # CentOS下查看PyTorch版本

2.3 自定义类处理

若模型包含自定义类，加载前需确保类定义已导入，否则会报错。

3. 实际应用示例

3.1 完整流程代码

# 训练并保存模型
model = MyModel()
train(model)
torch.save(model.state_dict(), 'model.pth')

# 加载模型
loaded_model = MyModel()
loaded_model.load_state_dict(torch.load('model.pth'))
loaded_model.eval()  # 切换到推理模式

3.2 多GPU训练模型的处理

若使用DataParallel训练，保存时需去掉模块前缀：

# 保存
torch.save(model.module.state_dict(), 'multi_gpu_model.pth')

# 加载
model = MyModel()
model = torch.nn.DataParallel(model)  # 重新包装
model.load_state_dict(torch.load('multi_gpu_model.pth'))