网络模型的保存与读取

方式1：

vgg16 = torchvision.models.vgg16(pretrained=False)
#保存方式1（模型结构+参数）
torch.save(vgg16, "vgg16_method1.pth")

#加载方式1
model = torch.load("D:\\Myproject\\pytorch-tutorial\\vgg16_method1.pth")
print(model)

方式2，字典形式：

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

终端输入dir:

-a----         2022/10/4     20:24      553448005 vgg16_method1.pth
-a----         2022/10/4     20:24      553440325 vgg16_method2.pth

方式二容量更小。

方式二转换成方式一：

#加载方式2
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("D:\\Myproject\\pytorch-tutorial\\vgg16_method2.pth"))
# model = torch.load("D:\\Myproject\\pytorch-tutorial\\vgg16_method2.pth")
print(vgg16)

陷阱：

# 陷阱
class Suibao(nn.Module):
    def __init__(self):
        super(Suibao, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3)

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


suibao = Suibao()
torch.save(suibao, "suibao_method1.pth")



# 陷阱1
from model_save import *
model = torch.load("D:\\Myproject\\pytorch-tutorial\\suibao_method1.pth")
print(model)