模型的保存与加载（pytorch）

（联邦学习笔记，资料来源于b站小土堆）

对于我们自定义的或者对于现存的模型进行修改后，我们可能要在别的地方进行调用，这时候需要先将它进行保存以便后续调用。

模型保存

对于模型的保存有两种方式，一种是同时保存模型结构和模型参数；另一种是仅保存模型参数（官方推荐）

保存方式一（保存模型结构和模型参数）:

torch.save(vgg16,"vgg16_methond1.pth")
           

其中vgg16是模型，vgg16_methond1.pth是保存的名称

保存方式二（仅保存模型参数）：

torch.save(vgg16.state_dict(),"vgg16_methond2.pth")
           

如果要保存自定义的模型，如下操作：

#保存自定义的模型
class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.conv1 = nn.Conv2d(3,64,kernel_size=3)

    def forward(self,input):
        output = self.conv1(input)
        return output

mymodule = MyModule()

torch.save(mymodule,"mymodule_save.pth")
           

以上是保存的相关操作，可定义在一个python文件中。完整代码如下：

import torch
import torchvision
from torch import nn

vgg16 = torchvision.models.vgg16(pretrained=False)


#保存方式一，保存模型结构和模型参数
torch.save(vgg16,"vgg16_methond1.pth")

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

#保存自定义的模型
class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.conv1 = nn.Conv2d(3,64,kernel_size=3)

    def forward(self,input):
        output = self.conv1(input)
        return output

mymodule = MyModule()

torch.save(mymodule,"mymodule_save.pth")
           

模型加载

加载内容一般根据的保存的方式来显示，一般形式为：model = torch.load("vgg16_methond1.pth")

对于上面的保存方式二，如果想要以（模型结构+模型参数）方式加载，则要先知道原本的模型是什么（上述是vgg16），然后进行操作，具体如下：

#加载方式2 同时加载模型和参数
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_methond2.pth"))
print(vgg16)
           

模型加载可写在同一个python文件中，完整代码如下：

import torch
import torchvision

import model_save


#加载以 方式一 保存的模型（模型结构和参数）
model = torch.load("vgg16_methond1.pth")
print(model)

#加载 方式二 保存的模型（模型参数）

#加载方式1 直接加载模型参数
model1 = torch.load("vgg16_methond2.pth")
print(model1)

#加载方式2 同时加载模型和参数
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_methond2.pth"))
print(vgg16)

#加载自定义模型
model2 = torch.load("mymodule_save.pth")
print(model2)
           

（如有不同意见，欢迎留下评论，我见到了会第一时间回复）