一 保存、加载自己的网络模型
保存:
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')