【pytorch】buffer的使用

写作背景：最近在研究pix2pix、cyclegan和pix2pixHD的代码，发现自己对pytorch根本就是什么都不会，所以准备开始针对代码中出现的自己不懂的地方写在CSDN上，记录自己的学习过程和学习难点，分享自己的学习心得。

buffer ：pytorch通常情况下，是将网络中的参数保存成OrderedDict形式的。其中的参数其实包括2种：一是模型中各种 module所包含的参数，即nn.Parameter，我们可以在网络中定义其他的nn.Parameter参数。另外一种是buffer。在使用的过程中，前者每次optim.step会得到更新，而不会更新后者。

class myModel(nn.Module):
    def __init__(self, kernel_size=3):
        super(Depth_guided1, self).__init__()
        self.kernel_size = kernel_size
        self.back_end = torch.nn.Sequential(
            torch.nn.Conv2d(3, 32, 3, padding=1),
            torch.nn.ReLU(True),
            torch.nn.Conv2d(3, 64, 3, padding=1),
            torch.nn.ReLU(True),
            torch.nn.Conv2d(64, 3, 3, padding=1),
            torch.nn.ReLU(True),
        )
 
        mybuffer = np.arange(1,10,1)
        self.mybuffer_tmp = np.randn((len(mybuffer), 1, 1, 10), dtype='float32')
        self.mybuffer_tmp = torch.from_numpy(self.mybuffer_tmp)
        # register preset variables as buffer
        # So that, in testing , we can use buffer variables.
        self.register_buffer('mybuffer', self.mybuffer_tmp)
 
        # Learnable weights
        self.conv_weights = nn.Parameter(torch.FloatTensor(64, 10).normal_(mean=0, std=0.01))
        # Other code
        def forward(self):
            ...
            # 这里使用 self.mybuffer!

通过代码，可以得出以下的一些小结论：

1. 定义Parameter和buffer都只需要传入Tensor即可。也不需要将其转成gpu。这是因为，当网络进行 .cuda()时候，会自动将里面的层的参数，buffer等转换成相应的GPU上；
2. self.register_buffer可以将tensor注册成buffer，在forward中使用self.mybuffer, 而不是self.mybuffer_tmp；
3. 网络存储时也会将buffer存下，当网络load模型时，会将存储的模型的buffer也进行赋值；
4. optim.step只能更新nn.Parameter类型的参数，而buffer的更新是在forward中。