pytorch(三)——保存自己的模型
模型依照前一片文章訓練好之後應該怎麼保存呢?
有兩個方法
1.保存模型參數
torch.save(the_model.state_dict(), PATH)
the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))
2.保存整個模型(推薦)
torch.save(the_model, PATH)
the_model = torch.load(PATH)
可能出現報錯
AttributeError: Can't get attribute 'Net' on module '__main__'
解決方法:補充上網絡的結構就可以了
測試
model = the_model()
print(model)
如果能打印出網絡結構,則成功了~~