加载和保存Pytorch模型

---

一 保存、加载自己的网络模型

保存：

net = LeNet(classes=2)
# 方式一：保存模型中的参数
path_model_state_dict = "../saved_models/model.pth"
torch.save(net.state_dict(), path_model_state_dict)

# 方式二：保存整个网络
path_model = "../saved_models/model.pkl"
torch.save(net, path_model)

加载：

# 方式一：加载模型中的参数
resnet=resnet50()    # 加载网络结构
path_model_state_dict = "../saved_models/model.pth" # 加载网络参数的路径
resnet.load_state_dict(torch.load(path_model_state_dict))  # 加载该网络结构的预训练参数成为模型
model_loaded = resnet  # 模型

# 方式二：加载整个网络
# 无需加载网络，因为模型保存了网络的所有信息
path_model="./models/model.pkl"
net_loaded=torch.load(path_model)

---

方式三：部分保存

---

二 加载官方模型

import torchvision.models as models
resnet50 = models.resnet50(pretrained=True)
# resnet50 = models.resnet50(pretrained=False)

保存为pt

net.eval()
trace = torch.jit.trace(net, torch.randn(1, 3, 512, 512))
torch.jit.save(trace,'UNet_model.pt')