pytorch读取与保存网络模型

pytorch读取与保存网络模型

方法一：

只保存网络模型的参数，而不保存网络的结构

保存：

torch.save(model.state_dict(), 'model_params.pth')

读取：

model.load_state_dict(torch.load('model_params.pth'))

方法二：

同时保持网络的结构和参数

保存：

torch.save(model, 'model.pth')

读取：

model = torch.load('model.pth')