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