pytorch（三）——保存自己的模型

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)

如果能打印出網絡結構，則成功了～～