模型的保存和加载

### 如何在 PyTorch 中保存和加载模型的最佳实践 #### 保存整个模型 为了同时保存模型的参数和网络结构信息，可以使用 `torch.save` 函数来序列化整个模型对象到磁盘。这方法简单直观，适合小型项目或快速原型开发。 ```python import torch # 假设 'the_model' 是已经训练好的模型实例 PATH = './model.pth' torch.save(the_model, PATH) # 保存模型[^1] ``` 这种方法的优点在于恢复时无需重新定义模型类即可直接加载并使用该模型： ```python # 加载已保存的模型 loaded_model = torch.load(PATH) loaded_model.eval() # 设置为评估模式 ``` 然而，在实际应用中更推荐的做法是仅保存模型的状态字典 (`state_dict`) 而不是完整的模型对象。这种方式更加灵活高效，并且能够更好地处理不同版本之间的兼容性问题。 #### 只保存状态字典 (State Dictionary) 对于只保存模型权重而不包括其架构的情况，则应采用如下方式操作: ```python torch.save(the_model.state_dict(), PATH) # 仅保存模型参数 ``` 当需要再次部署此模型时，先创建一个新的相同类型的未初始化模型实例，再将其状态字典更新为目标文件中的数据： ```python new_model = TheModelClass(*args, **kwargs) # 创建新的模型实例 new_model.load_state_dict(torch.load(PATH)) # 将之前保存的参数赋给新模型 new_model.eval() ``` 这种做法不仅减少了存储空间占用量，还便于迁移学习以及微调预训练过的神经网络。 #### 使用 CUDA 进行加速 如果希望利用 GPU 来加快计算速度，可以在构建模型前设置环境变量 `cudnn.benchmark=True` ，从而让 cuDNN 自动寻找最适合当前硬件配置下的卷积算法实现方案[^2]。 ```python import torch.backends.cudnn as cudnn cudnn.benchmark = True device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') model.to(device) ``` 通过上述措施可有效提升含有大量重复运算任务（如图像分类、目标检测等）的应用程序性能表现。