Pytorch：保存模型或 state_dict 给出不同的磁盘空间占用

Posted 2023-03-12

【中文标题】Pytorch：保存模型或 state_dict 给出不同的磁盘空间占用

【英文标题】：Pytorch: saving model or state_dict gives different on-disk space occupation

【发布时间】：2019-12-15 21:21:37

【问题描述】：

我在玩torch.save 函数时发现了一些奇怪的东西，假设我从torchvision 存储库加载了一个模型：

model = torchvision.models.mobilenet_v2()

如果我以这种方式保存模型：

torch.save(model,'model.pth')

我得到一个 14MB 的文件，如果我这样做：

torch.save(model.state_dict(),'state_dict.pth')

文件大小增加到约 500MB。由于我没有找到有关此行为的任何参考，我想知道是什么导致了大小的增加。它与压缩有关吗？保存整个state_dict 是否会存储额外的东西，比如未初始化的渐变？

附： vgg16等其他模型也会发生同样的情况

【问题讨论】：

Best way to save a trained model in PyTorch?的可能重复

@dennlinger 在问这个问题之前我已经看到了这个问题，它没有给出任何提示。

PyTorch: What's the difference between state_dict and parameters()? 提供了更多见解，但两者之间的差异很大

你的 PyTorch 版本是多少？

【参考方案1】：

如果询问模型中的内容：

vars(vgg16)

输出：

'_backend': <torch.nn.backends.thnn.THNNFunctionBackend at 0x232c78759b0>,
 '_parameters': OrderedDict(),
 '_buffers': OrderedDict(),
 '_backward_hooks': OrderedDict(),
 '_forward_hooks': OrderedDict(),
 '_forward_pre_hooks': OrderedDict(),
 '_state_dict_hooks': OrderedDict(),
 '_load_state_dict_pre_hooks': OrderedDict(),
 '_modules': OrderedDict([('features', Sequential(
                 (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                 (1): ReLU(inplace)
                 (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                 (3): ReLU(inplace)
                 (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
                 (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                 (6): ReLU(inplace)
                 (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                 (8): ReLU(inplace)
                 (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
                 (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                 (11): ReLU(inplace)
                 (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                 (13): ReLU(inplace)
                 (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                 (15): ReLU(inplace)
                 (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
                 (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                 (18): ReLU(inplace)
                 (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                 (20): ReLU(inplace)
                 (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                 (22): ReLU(inplace)
                 (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
                 (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                 (25): ReLU(inplace)
                 (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                 (27): ReLU(inplace)
                 (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
                 (29): ReLU(inplace)
                 (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
               )),
              ('avgpool', AdaptiveAvgPool2d(output_size=(7, 7))),
              ('classifier', Sequential(
                 (0): Linear(in_features=25088, out_features=4096, bias=True)
                 (1): ReLU(inplace)
                 (2): Dropout(p=0.5)
                 (3): Linear(in_features=4096, out_features=4096, bias=True)
                 (4): ReLU(inplace)
                 (5): Dropout(p=0.5)
                 (6): Linear(in_features=4096, out_features=1000, bias=True)
               ))]),
 'training': True

你会得到它不仅仅是状态字典。

vgg16.state_dict()

状态字典在 _modules (vgg16._modules['features'].state_dict()) 内

这就是为什么当您保存模型时，您不仅保存状态字典，还保存所有上述内容，例如参数、缓冲区、挂钩...

但是，如果您不使用参数、缓冲区、挂钩来为您的模型进行推理时间，您可能会避免保存这些。

---

保存时的尺寸：

torch.save(model,'model.pth')
torch.save(model.state_dict(),'state_dict.pth')

应该是：model.pth > state_dict.pth，因为状态字典包含在模型中。

【参考方案2】：

嘿，我删除了我的最后一个答案，因为我错了，截至 pytorch 版本 1.1.0a0+863818e，使用：

torch.save(model,'model.pth')

并使用：

torch.save(model.state_dict(),'state_dict.pth')

在大小上给出了相同的结果，您使用的是正确加载网络吗？证明：

-rw-rw-r-- 1 bpinaya bpinaya  14M Aug  8 10:26 model.pth
-rw-rw-r-- 1 bpinaya bpinaya  14M Aug  8 10:27 state_dict.pth
-rw-rw-r-- 1 bpinaya bpinaya 528M Aug  8 10:29 vgg.pth
-rw-rw-r-- 1 bpinaya bpinaya 528M Aug  8 10:29 vggstate_dict.pth

【讨论】：

我在保存模型之前一直在做一些forwards，可能缓冲区还在，但我真的在字典中看不到它们

你原来的答案没有错。实际上保存模型和 model.state_dict() 会有不同的大小，当模型很大时会略有不同。 ;)