114_PyTorch 进阶：模型保存与读取的两大方式及“陷阱”避坑指南

在 PyTorch 中保存和加载模型主要有两种方式。虽然都能达到目的但它们在存储内容和使用便捷性上有着本质的区别。本文将结合实战代码带你掌握这些核心技巧。1. 方式一保存整个网络模型结构与参数这种方式最为直观它会把模型的结构定义和权重参数全部打包保存到一个文件中。代码实现Pythonimport torch import torchvision vgg16 torchvision.models.vgg16(pretrainedFalse) # 保存方式一保存模型结构模型参数 torch.save(vgg16, vgg16_method1.pth) # 加载方式一 model torch.load(vgg16_method1.pth) print(model)优点加载方便一行代码搞定。缺点文件体积较大存在“陷阱”见下文。2. 方式二只保存模型参数官方推荐这种方式只保存模型中每个层的权重和偏置即state_dict。加载时需要先创建一个模型实例再将参数“填”进去。代码实现Python# 保存方式二将模型参数保存为字典格式官方推荐体积小 torch.save(vgg16.state_dict(), vgg16_method2.pth) # 加载方式二 # 1. 必须先新建一个网络模型实例 vgg16_new torchvision.models.vgg16(pretrainedFalse) # 2. 将保存的参数字典加载到模型中 vgg16_new.load_state_dict(torch.load(vgg16_method2.pth)) print(vgg16_new)优点更安全、更节省空间是工业界和官方推荐的标准做法。3. 核心避坑指南自定义模型的加载“陷阱”文件特别提到了一个新手常犯的错误当你用“方式一”加载自定义模型时如果当前环境没有该模型的类定义程序会报错。错误场景如果你在A.py中定义并保存了模型在B.py中直接torch.load会提示AttributeError。解决方案在加载代码前必须确保模型类的定义是可访问的。Pythonimport torch # 陷阱必须把模型的定义复制过来或者从其他文件 import 进来 class Tudui(torch.nn.Module): def __init__(self): super(Tudui, self).__init__() self.conv1 torch.nn.Conv2d(3, 64, 3) def forward(self, x): return self.conv1(x) # 这样才能成功加载方式一保存的文件 model torch.load(tudui_method1.pth) print(model)或者使用更优雅的导入方式Pythonfrom model_save import * # 确保 Tudui 类被正确引入 model torch.load(tudui_method1.pth)4. 总结如何选择保存方式方式一适合快速实验想省去重新定义结构的麻烦但要记住加载时引入类定义。方式二适合生产环境不仅体积小且能更灵活地将参数加载到结构略有差异的模型中。 学习小结掌握了模型的保存与读取你就真正拥有了“复用”训练成果的能力。无论是模型的版本控制还是将模型部署到云端或移动端这都是必经之路。