Pytorch模型中的parameter與buffer（torch.nn.Module的成員）

前言：

class DaGMM(nn.Module):  # 自定義的模型需要繼承nn.Module(固定寫法)
    """Residual Block(殘塊)."""

    def __init__(self, n_gmm=2, latent_dim=3):
        super(DaGMM, self).__init__()  # (固定寫法)

        layers = []
        layers += [nn.Linear(118, 60)]
        layers += [nn.Tanh()]  # 激活函數
        layers += [nn.Linear(60, 30)]
        layers += [nn.Tanh()]
        layers += [nn.Linear(30, 10)]
        layers += [nn.Tanh()]
        layers += [nn.Linear(10, 1)]

        self.encoder = nn.Sequential(*layers)

        layers = []
        layers += [nn.Linear(1, 10)]
        layers += [nn.Tanh()]
        layers += [nn.Linear(10, 30)]
        layers += [nn.Tanh()]
        layers += [nn.Linear(30, 60)]
        layers += [nn.Tanh()]
        layers += [nn.Linear(60, 118)]

        self.decoder = nn.Sequential(*layers)

        layers = []
        layers += [nn.Linear(latent_dim, 10)]
        layers += [nn.Tanh()]
        layers += [nn.Dropout(p=0.5)]
        layers += [nn.Linear(10, n_gmm)]
        layers += [nn.Softmax(dim=1)]

        self.estimation = nn.Sequential(*layers)

38         self.register_buffer("phi", torch.zeros(n_gmm))
39         self.register_buffer("mu", torch.zeros(n_gmm, latent_dim))
40         self.register_buffer("cov", torch.zeros(n_gmm, latent_dim, latent_dim))

　　我們知道，pytorch一般情況下，是將網絡中的參數保存成OrderedDict（見附1）形式的。這里的參數其實包括2種：一種是模型中的各種module含的參數，即nn.Parameter，我們當然可以在網絡中定義其他的nn.Parameter參數。另外一種是buffer。前者每次optim.step會得到更新，而不會更新后者。

1、模型保存

　在Pytorch中一種模型保存和加載的方式如下：

# save
torch.save(model.state_dict(), PATH)

# load
model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

　　可以看到，模型保存的是model.state.dict()的返回對象。model.state_dict的返回對象是一個OrderedDict，它以鍵值對的形式保存模型中需要保存下來的參數，例如：

import torch
import torch.nn as nn

class MyModule(nn.Module):
    def __init__(self, input_size, output_size):
        super(MyModule, self).__init__()
        self.lin = nn.Linear(input_size, output_size)

    def forward(self, x):
        return self.lin(x)

module = MyModule(4, 2)
print(module.state_dict())
輸出：
OrderedDict([('lin.weight', tensor([[-0.3636, -0.4864, -0.2716, -0.4416],
        [ 0.0119,  0.4462, -0.4558,  0.0188]])), ('lin.bias', tensor([ 0.1056, -0.1058]))])

　　模型中的參數就是線性層的weight和bias。

2、Parameter & buffer

　torch.nn.register_parameter()用於注冊Parameter實例到當前Module中（一般可以用torch.nn.Parameter()代替）；torch.nn.register_buffer()用於注冊Buffer實例到當前Module中。此外，Module中的parameters()函數會返回當前Module中所注冊的所有Parameter的迭代器；而_all_buffers()函數會返回當前Module中所注冊的所有Buffer的迭代器，（所以優化器不會計算Buffer的梯度，自然不會對其更新）。此外，Module中的state_dict()會返回包含當前Module中所注冊的所有Parameter和Buffer（所以模型中未注冊成Parameter或Buffer的參數無法被保存）。

　模型中需要保存下來的參數包括兩種：

• 一種是反向傳播需要被optimizer更新的，稱之為parameter
• 一種是反向傳播不需要被optimizer更新的，稱之為buffer

　第一種參數我們可以通過model.parameter()返回；第二種參數我們可以通過model.buffers()返回。因為我們模型保存的是state_dict返回的OrderedDict，所以這兩種參數不僅要滿足是/否需要更新的要求，還需要被保存到OrderedDict。

　那么現在的問題是這兩種參數如何創建呢？創建好之后如何保存到OrderedDict呢？

　2.1、parameter參數有兩種創建方式：　

 　　1、我們可以直接將模型的成員變量（self.xxx）通過nn.Parameter()創建，會自動注冊到parameters中，可以通過model.parameters()返回，並且這樣創建的參數會自動保存到OrderedDict中去。

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.my_param = nn.Parameter(torch.randn(3, 3))  # 模型的成員變量
    def forward(self, x):
        # 可以通過 self.my_param 和 self.my_buffer 訪問
        pass

model = MyModel()
for param in model.parameters():
    print(param)
print("----------------")
print(model.state_dict())
輸出：
Parameter containing:
tensor([[-0.5421,  2.9562,  0.3447],
        [ 0.0869, -0.3464,  1.1299],
        [ 0.8644, -0.1384, -0.6338]])
----------------
OrderedDict([('param', tensor([[-0.5421,  2.9562,  0.3447],
        [ 0.0869, -0.3464,  1.1299],
        [ 0.8644, -0.1384, -0.6338]]))])

　　2、通過nn.Parameter()創建普通的Parameter對象，不作為模型的成員變量，然后將Parameter對象通過register_parameter()進行注冊，可以通過model.parameters()返回，注冊后的參數也是會自動保存到OrderedDict中去。

import torch
import torch.nn as nn
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

        param = nn.Parameter(torch.randn(3, 3))  # 普通 Parameter 對象
        self.register_parameter("my_param", param)

    def forward(self, x):
        # 可以通過 self.my_param 和 self.my_buffer 訪問
        pass
model = MyModel()
for param in model.parameters():
    print(param)
print("----------------")
print(model.state_dict())
輸出：
Parameter containing:
tensor([[-0.2313, -0.1490, -1.3148],
        [-1.2862, -2.2740,  1.0558],
        [-0.6559,  0.4552,  0.5993]])
----------------
OrderedDict([('my_param', tensor([[-0.2313, -0.1490, -1.3148],
        [-1.2862, -2.2740,  1.0558],
        [-0.6559,  0.4552,  0.5993]]))])

 　2.2、buffer參數的創建方式：

　　這種參數的創建需要先創建tensor，然后將tensor通過register_buffer()進行注冊，可以通過model._all_buffers()返回，注冊完成后參數也會自動保存到OrderedDict中去。

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

        buffer = torch.randn(2, 3)  # tensor
        self.register_buffer('my_buffer', buffer)

    def forward(self, x):
        # 可以通過 self.param 和 self.my_buffer 訪問
        pass
model = MyModel()
for buffer in model._all_buffers():
    print(buffer)
print("----------------")
print(model.state_dict())
輸出：
tensor([[-0.2191,  0.1378, -1.5544],
        [-0.4343,  0.1329, -0.3834]])
----------------
OrderedDict([('my_buffer', tensor([[-0.2191,  0.1378, -1.5544],
        [-0.4343,  0.1329, -0.3834]]))])

總結：

　I、模型中需要進行更新的參數注冊為Parameter，不需要進行更新的參數注冊為buffer

　II、模型保存的參數是Model.state_dict()返回的OrderedDict

　III、模型進行設備移動時（CPU--->GPU），模型中注冊的參數（Parameter和buffer）會同時進行移動。

附錄：

1、很多人認為python中的字典是無序的，因為它是按照hash來存儲的，但是python中有個模塊collection，里面自帶了一個子類OrderedDict，實現了對字典對象中元素的排序。

import collections

print("Regular dictionary")
d = {}
d['a'] = 'A'
d['b'] = 'B'
d['c'] = 'C'
for k, v in d.items():
    print(k, v)

print("\nOrder dictionary")
d1 = collections.OrderedDict()
d1['a'] = 'A'
d1['b'] = 'B'
d1['c'] = 'C'
d1['1'] = '1'
d1['2'] = '2'
for k, v in d1.items():
    print(k, v)
輸出：
Regular dictionary
a A
c C
b B

Order dictionary
a A
b B
c C
1 1
2 2

可以看到，同樣是保存了ABC等幾個元素，但是使用OrderedDict會根據放入元素的先后順序進行排序。所以輸出的值是拍好序的，OrderedDict對象的字典對象，如果其順序不同，那么python會把他們當做兩個不同的對象。

參考：

參考1：【Pytorch】模型中buffer的使用

參考2：python中OrdredDict用法

參考3：Pytorch模型中的parameter和buffer

參考4：Pytorch采坑記錄