Pytorch两种模型保存方式

Pytorch两种模型保存方式

大黑_7e1b 关注

2019.02.12 17:49:35 字数 13 阅读 5,907

只保存模型参数

# 保存
torch.save(model.state_dict(), '\parameter.pkl')
# 加载
model = TheModelClass(...)
model.load_state_dict(torch.load('\parameter.pkl'))

保存完整模型

# 保存
torch.save(model, '\model.pkl')
# 加载
model = torch.load('\model.pkl')

1人点赞

学习笔记

"小礼物走一走，来简书关注我"

赞赏支持 还没有人赞赏，支持一下

大黑_7e1b

总资产2 (约0.20元) 共写了1058字 获得7个赞 共3个粉丝

关注