保存
import torch
import torchvision
vgg16 = torchvision.models.vgg16(pretrained=False)
# 保存方式1 保存模型结构 + 模型参数
torch.save(vgg16, "vgg16_method1.pth")
# 保存方式2 保存模型参数(官方推荐)
torch.save(vgg16.state_dict(), "vgg16_method2.pth")
加载
import torch
import torchvision.models
from torch import nn
# 方式1,加载模型
# model = torch.load("vgg16_method1.pth")
# print(model)
# 方式2, 加载模型
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))
# print(vgg16)