Ubuntu18.04LTS日志10-pytorch训练模型与保存

最新推荐文章于 2023-06-01 10:55:59 发布

会飞的糖企鹅

最新推荐文章于 2023-06-01 10:55:59 发布

阅读量354

点赞数

分类专栏： Docker Pytorch

本文链接：https://blog.csdn.net/weixin_41666001/article/details/103028592

版权

Docker 同时被 2 个专栏收录

7 篇文章 0 订阅

订阅专栏

Pytorch

6 篇文章 0 订阅

订阅专栏

Ubuntu18.04LTS日志10-pytorch训练模型与保存

参考链接：
https://blog.csdn.net/sll71/article/details/95856952

//  save model
torch.save(model, './model/model.pkl')  # save entire net
torch.save(model.state_dict(), './model/model_params.pkl')  # save only the parameters

//load model
model = torch.load('/root/data/Tom/dogvscatv01/model/model.pkl')  # 加载模型
if torch.cuda.is_available():
  print('yes')
  device = torch.device("cuda")
  model = model.to(device)