pytorch：保存模型或state_dict提供了不同的磁盤空間占用

Question

我正在使用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存儲未初始化的漸變之類的額外內容？

PS對於類似vgg16其他模型vgg16

Answer 1

嘿，自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

Answer 2

如果詢問模型是什么：

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()

狀態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因為狀態dict包含在模型中。