Pytorch中保存模型的两种方法
Pytorch中保存模型的两种方法
·

💪 专业从事且热爱图像处理,图像处理专栏更新如下👇:
📝《图像去噪》
📝《超分辨率重建》
📝《语义分割》
📝《风格迁移》
📝《目标检测》
📝《图像增强》
📝《模型优化》
📝《模型实战部署》
📝《图像配准融合》
📝《数据集》
📝《高效助手》

一、保存整个模型
这种方法会保存模型的整个定义(包括网络结构和参数)。这样保存的模型可以直接被加载并使用,但前提是保存和加载模型的代码环境需要完全一致。
保存模型时使用代码:
torch.save(model, 'model.pth')
加载模型时使用代码:
model = torch.load('model.pth')
二、只保存模型参数
这种方法只会保存模型的状态字典(state_dict),即模型的参数。加载这样的模型需要先定义模型的结构,然后加载参数。
保存模型使用代码:
torch.save(model.state_dict(), 'params.pth')
加载模型使用代码:
model = Model() # 需要先定义模型结构
model.load_state_dict(torch.load('params.pth'))
三、总结
以上就是Pytorch中保存模型的两种方法,学者在训练前确定好保存方法,推理测试阶段对应的调整代码。
感谢您阅读到最后!😊总结不易,多多支持呀🌹 点赞👍收藏⭐评论✍️,您的三连是我持续更新的动力💖
关注公众号「视觉研坊」,获取干货教程、实战案例、技术解答、行业资讯!
更多推荐


所有评论(0)