训练模型时保存/读取模型参数

保存模型参数

# 保存方式1,模型结构+模型参数
torch.save(vgg16, "vgg16_method1.pth")
# 保存方式2,模型参数（官方推荐）
torch.save(vgg16.state_dict(), "vgg16_method2.pth")

读取模型参数

# 读取方式1-->保存方式1
model1 = torch.load("vgg16_method1.pth")
print(model1)

# 读取方式2-->保存方式2
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))
print(vgg16)