PyTorch保存和加载模型

Posted 2021-02-20 picassooo

在PyTorch中使用torch.save来保存模型的结构和参数，有两种保存方式：

# 方式一：保存模型的结果信息和参数信息
torch.save(model, ‘./model.pth‘)

# 方式二：仅保存模型的参数信息
torch.save(model.state_dict(), ‘./model_state.pth‘)

相应的，有两种加载模型的方式：

# 方式一：加载完整的模型结构和参数信息，在网络较大时加载时间比较长，同时存储空间也比较大
model1= torch.load(‘model.pth‘)   

# 方式二：需先搭建网络模型model2，然后通过下面的语句加载参数
model2.load_state_dic(torch.load(‘model_state.pth‘))