开始使用免费开始使用

保存并加载模型

一家制造企业希望基于图像对其项目进行分类,并确定合适的运输包装。您已经在 PyTorch 中训练了一个高精度模型,现在计划将该模型及其预训练权重保存下来,以便后续使用并与团队共享,确保他们可以顺利加载。

torchtorch.nn(别名为 nn)已导入。预训练的模型对象已在您的工作区中,名称为 model,其架构类为 ManufacturingCNN

本练习是课程的一部分

使用 PyTorch 进行图像深度学习

查看课程

练习说明

  • 将预训练模型保存为 ModelCNN.pth,注意只保存权重,不仅是架构。
  • 使用类 ManufacturingCNN() 创建一个名为 loaded_model 的模型实例。
  • ModelCNN.pth 中的权重加载到 loaded_model,方法是将权重传入 .load_state_dict()

交互式实操练习

通过完成这段示例代码来试试这个练习。

# Save the model
torch.____(model.____, ____)

# Create a new model
loaded_model = ____

# Load the saved model
loaded_model.____(torch.____('ModelCNN.pth'))
print(loaded_model)
编辑并运行代码