PyTorch 保存模型结构参数及加载模型

PyTorch 保存模型结构参数及加载模型

保存模型与加载

保存模型分为两种方式：

1. 保存整个网络结构和参数
2. 保存整个网络的参数

# 1.保存并加载整个网络结构和参数
# 保存模型
torch.save(net, 'net.pth')
# 加载模型
net = torch.load('net.pth')

# 2.保存并加载整个网络的参数
# 保存模型
torch.save(net.state_dict(), 'net.pth')
# 加载模型
net.load_state_dict(torch.load('net.pth'))

Ref

1. https://pytorch.org/tutorials/beginner/saving_loading_models.html
2. PyTorch学习：加载模型和参数