Pytorch兩種模型保存方式

只保存模型參數(shù)

# 保存
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')

總結(jié)

以上是生活随笔為你收集整理的Pytorch两种模型保存方式的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

如果覺(jué)得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。