pytorch 保存和加载模型

Posted 2023-02-23

【中文标题】pytorch 保存和加载模型

【英文标题】：pytorch save and load model

【发布时间】：2021-12-15 14:45:00

【问题描述】：

原始模型和保存后加载模型有什么区别吗？

在训练之前，我只是保存模型然后加载，因为我想知道在保存和加载过程中是否有任何变化。

这是我的代码

只是用于测试的模型

class test_model(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1 = nn.Conv2d(in_channels = 3, out_channels = 50, kernel_size = 3, stride=1, padding=1, bias = False)
        self.maxpool1 = nn.MaxPool2d(2, 2)
        self.bn1 = nn.BatchNorm2d(50)
        
        self.conv2_ = nn.Conv2d(in_channels = 50, out_channels = 10, kernel_size = 1, stride=1, padding=0, bias = False)
        self.conv2 = nn.Conv2d(in_channels = 10, out_channels = 50, kernel_size = 3, stride=1, padding=1, bias = False)
        self.maxpool2 = nn.MaxPool2d(2, 2)
        self.bn2 = nn.BatchNorm2d(50)

        self.conv3_ = nn.Conv2d(in_channels = 50, out_channels = 10, kernel_size = 1, stride=1, padding=0, bias = False)
        self.conv3 = nn.Conv2d(in_channels = 10, out_channels = 50, kernel_size = 3, stride=1, padding=1, bias = False)
        self.maxpool3 = nn.MaxPool2d(2, 2)
        self.bn3 = nn.BatchNorm2d(50)
        
        self.conv4_ = nn.Conv2d(in_channels = 50, out_channels = 20, kernel_size = 1, stride=1, padding=0, bias = False)
        self.conv4 = nn.Conv2d(in_channels =20, out_channels = 100, kernel_size = 3, stride=1, padding=1, bias = False)
        self.maxpool4 = nn.MaxPool2d(2, 2)
        self.bn4 = nn.BatchNorm2d(100)

        self.conv5_ = nn.Conv2d(in_channels = 100, out_channels = 10, kernel_size = 1, stride=1, padding=0, bias = False)
        self.conv5 = nn.Conv2d(in_channels = 10, out_channels = 100, kernel_size = 3, stride=1, padding=1, bias = False)
        self.maxpool5 = nn.MaxPool2d(2, 2)
        self.bn5 = nn.BatchNorm2d(100)
        
        self.fc = nn.Sequential(Flatten(),
                               nn.Linear(100*7*7, 100),
                                nn.ReLU(),
                               nn.Linear(100,100))
        
    def forward(self, inputs):
        feature_map1 = self.conv1(inputs)
        feature_map1 = self.maxpool1(feature_map1)
        feature_map1 = self.bn1(feature_map1)
        
        feature_map2 = self.conv2_(feature_map1)
        feature_map2 = self.conv2(feature_map2)
        feature_map2 = self.maxpool2(feature_map2)
        feature_map2 = self.bn2(feature_map2)
         
        feature_map3 = self.conv3_(feature_map2)   
        feature_map3 = self.conv3(feature_map3)
        feature_map3 = self.maxpool3(feature_map3)
        feature_map3 = self.bn3(feature_map3)
        
        feature_map4 = self.conv4_(feature_map3)
        feature_map4 = self.conv4(feature_map4)
        feature_map4 = self.maxpool4(feature_map4)
        feature_map4 = self.bn4(feature_map4)
        
        feature_map5 = self.conv5_(feature_map4)
        feature_map5 = self.conv5(feature_map5)
        feature_map5 = self.maxpool5(feature_map5)
        feature_map5 = self.bn5(feature_map5)
           
        output = self.fc(feature_map5)
        
        return output

然后

model_cpu = test_model()

保存和加载

torch.save(model_cpu, '/home/mskang/hyeokjong/model_cpu.pt')
model_load = torch.load('/home/mskang/hyeokjong/model_cpu.pt')

和

model_load  == model_cpu
------------------------------------
False

然而

print(model_load)
print(model_cpu)

好像都一样

还有

我还训练了两个模型（model_load，model_cpu），结果看起来也一样。

所以我认为那些拖车模型是相同的，应该是相同的 但是为什么False

【参考方案1】：

它们具有相同的底层模型，但是是不同的 Python 对象。这就是为什么__eq__ 在尝试model_load == model_cpu 时返回False。您可以将model_load 和model_cpu 视为相同nn.Module 的两个副本。