0
点赞
收藏
分享

微信扫一扫

[pytorch][原创]保存模型和权重

探头的新芽 2024-11-04 阅读 9

只保存模型参数
 # 保存
 torch.save(model.state_dict(), '\model.pth')
 # 加载
 model = TheModelClass(...)
 model.load_state_dict(torch.load('\model.pth'))
 保存完整模型
 # 保存
 torch.save(model, '\model.pth')
 # 加载
 model = torch.load('\model.pth')

举报

相关推荐

0 条评论