当前位置: 代码迷 >> 综合 >> [pytorch][原创]保存模型和权重
  详细解决方案

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

热度:81   发布时间:2024-03-09 06:27:33.0

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