模型保存和修改

介绍

就是把参数保存下来。

保存的文件实际上没啥要求的，但是建议后缀.pth

两种方式保存（1、存模型+参数 2、参数存为字典）

vgg16 = torchvision.models.vgg16(pretrained=False)
# 保存方式1,模型结构+模型参数   模型 + 参数 都保存
torch.save(vgg16, "vgg16_method1.pth")  # 引号里是保存路径
# 保存方式2，模型参数（官方推荐） ，因为这个方式，储存量小，在terminal中，ls -all可以查看
 torch.save(vgg16.state_dict(), "vgg16_method2.pth")

对应加载方式

# 方式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"))
model = torch.load("vgg16_method2.pth")
print(vgg16)

陷阱

即方式1需要声明下模型定义/导入下模型类。

# 陷阱，用第一种方式保存时，如果是自己的模型，就需要在加载中，把class重新写一遍，但并不需要实例化，即可
# 这个陷阱，也是可以避免的，最上面的 from model_save import *，就是在做这个事情，避免出现错误
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3)

    def forward(self, x):
        x = self.conv1(x)
        return x

model = torch.load('tudui_method1.pth')
print(model)