深入理解Pytorch之register_buffer

使用

import torch.nn as nn
import torch
class net(nn.Module):
    def __init__(self):
        super(net,self).__init__()
        self.register_buffer("a",torch.ones(2,3))#从此，self.a其实就是torch.ones(2,3)。
    def forward(self,x):
        return x+self.a#使用

理解

register_buffer的作用是将torch.ones(2,3)这个tensor注册到模型的 buffers() 属性中，并命名为a,这代表a对应的是一个持久态，不会有梯度传播给它，但是能被模型的state_dict记录下来。可以理解为模型的常数。

注意，没有保存到模型的 buffers() 或 parameters() 属性中的参数是不会被记录到state_dict中的，在 buffers() 中的参数默认不会有梯度，parameters() 中的则相反。

我们可以将前者理解为常数，后者理解为变量。

1. requires_grad=False
2. 不会注册到模型参数中model.parameters()
3. 会注册到模型model.state_dict()中。

一个很多人疑问的问题是：既然register_buffer的对象是模型中的常数，那为什么不直接使用下面的方法一，还不更直接吗？

class net(nn.Module):
    def __init__(self,x=None):
        super(net,self).__init__()
		self.a=torch.ones(2,3)#方法一
		self.register_buffer("a",torch.ones(2,3))#方法二

这么跟你说吧，如果常数是这种torch.ones(2,3)的话，两者确实在使用体验上没有任何差别（虽然后者会把torch.ones(2,3)这个常数注册到model.state_dict()中，前者不会）。

但是，我们可能会遇到这样的场景：那个常数不是这么简单的常数，而是外部传入的。

class net(nn.Module):
    def __init__(self,x=None):
        super(net,self).__init__()
		self.a=x#方法一
		self.register_buffer("a",x)#方法二

x=**
x=***
x=**
#第一次运行的时候，你经过千辛万苦得到了模型中的常数x。
model=net(x)
#训练模型
#保存模型。
#完毕

#如果是方法一，你又要运行一遍获得x的过程。
x=**
x=***
x=**
model=net(x)
#载入模型model.load
#使用模型

#如果是方法二，不需要获得x，因为register_buffer会将常数x保存在state_dict中，载入就行了。
model=net(x)
#载入模型model.load
#使用模型