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

阅读 18

2024-11-04

只保存模型参数
 # 保存
 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)

0 0 举报