What is the difference between `register_buffer` and `register_parameter` of `nn.Module`

coincheung: <p>I was reading the code of mask-rcnn to see how they fix their bn parameters. I notice that they use <code>self.register_buffer</code> to create the <code>weight</code> and <code>bias</code>, while, in the pytorch BN definition,  <code>self.register_parameter</code> is used when <code>affine=True</code>. Could I simply think that <code>buffer</code> and <code>parameter</code> have everything in common except that <code>buffer</code> will neglect the operations to compute grad and update its values ?</p>
<p>By the way, what is the different between directly defining a <code>nn.Paramter</code> in the module and using <code>register_parameter</code> ?</p>

ptrblck: <p>Yes, you are correct in your assumption. If you have parameters in your model, which should be saved and restored in the <code>state_dict</code>, but not trained by the optimizer, you should register them as buffers.<br>
Buffers won’t be returned in <code>model.parameters()</code>, so that the optimizer won’t have a change to update them.</p>
<p>Both approaches work the same regarding training etc.<br>
There are some differences in the function calls however. Using <code>register_parameter</code> you have to pass the name as a string, which can make the creation of a range of parameters convenient. Besides that I think it’s just coding style which one you prefer.</p>

pechyonkin: <p>If I have some parameters that I don’t want to be trained, can I just add them as <code>self.some_params</code> inside the <code>nn.Module</code> to preserve state? Does <code>register_buffer</code> do anything special in that case as compared to just storing it inside <code>self</code>?</p>

pierrecurie: <p>@ <a href="https://discuss.pytorch.org/u/pechyonkin">pechyonkin</a></p>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="32723">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>If you have parameters in your model, which should be saved and restored in the <code>state_dict</code> , but not trained by the optimizer, you should register them as buffers.</p>
</blockquote>
</aside>
<p>If your self.some_params are nn.Parameter objects, then you don’t have to worry about this.  If they’re tensors, then they won’t be in the state_dict (unless registered as buffer).</p>

pechyonkin: <p>What are the downsides of not using a buffer? I am currently using <code>self.some_param</code> inside <code>nn.Module</code> to keep a tensor that keeps track of running average statistics of activations. I don’t need it for backprop, only to make decisions during runtime. I want to learn more about why my approach is not an optimal one. If you could explain or give some readings, that’d be great.</p>

pierrecurie: <p>Do you want it in the state_dict?</p>

pechyonkin: <p>I am sorry if this is a stupid question, but I am not sure if I want that. I checked <a href="https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict" rel="nofollow noopener">this</a>, but I still don’t see why I would need that. Would I need buffers if I want to save the model later? Are there any other reasons I would like to use <code>state_dict</code> rather than just assigning to <code>self</code>?</p>

ptrblck: <p>As <a class="mention" href="/u/pierrecurie">@pierrecurie</a> explained, one  reason to register the tensor as a buffer is to be able to serialize the model and restore all internal states.<br>
Another one is that all buffers and parameters will be pushed to the device, if called on the parent model:</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.my_tensor = torch.randn(1)
        self.register_buffer('my_buffer', torch.randn(1))
        self.my_param = nn.Parameter(torch.randn(1))
        
    def forward(self, x):
            return x

model = MyModel()
print(model.my_tensor)
&gt; tensor([0.9329])
print(model.state_dict())
&gt; OrderedDict([('my_param', tensor([-0.2471])), ('my_buffer', tensor([1.2112]))])

model.cuda()
print(model.my_tensor)
&gt; tensor([0.9329])
print(model.state_dict())
&gt; OrderedDict([('my_param', tensor([-0.2471], device='cuda:0')), ('my_buffer', tensor([1.2112], device='cuda:0'))])
</code></pre>
<p>As you can see, <code>model.my_tensor</code> is still on the CPU, where is was created, while all parameters and buffers were pushed to the GPU after calling <code>model.cuda()</code>.</p>

pechyonkin: <p>Thanks for clarification! Now it makes total sense. I will actually use buffer since I am going to use GPU at some point for the model I am building.</p>

oli4100: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> probably another dumb question, but why wouldn’t I just use nn.Parameter for both my_tensor and my_param and just state ‘requires_grad=False’ for the first? How would that be different to the example in your post?</p>

ptrblck: <p>I think there wouldn’t be a difference regarding the model training, gradient flow etc., so you could probably use this approach.<br>
However, it might be confusing to other users who are using your code to see some “buffers” in <code>model.parameters()</code>. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"><br>
Also, you would pass these buffers to the optimizer, if you just pass all <code>model.parameters()</code>.<br>
Again, this won’t mess with your training, but the optimizer will unnecessarily have to skip these buffers in its <code>step()</code> method.</p>
<p>I would describe it as a “clean” code style to separate buffers and parameters.</p>

oli4100: <p>Ah, thanks. An example where I find this distinction difficult is in the context of fixed positional encodings in the Transformer model. Typically I see implementations where the fixed positional encodings are registered as buffers but I’d consider these tensors as non-learnable parameters (that should show up in the list of model parameters), especially when comparing between methods that don’t rely on such injection of fixed tensors.</p>
<p>Re. your last remark, I guess <a href="https://github.com/pytorch/pytorch/issues/679" rel="nofollow noopener">this</a> should do the trick, but from that thread I understand it is poor coding practice.</p>
<p>So in general:<br>
buffers = ‘fixed tensors / non-learnable parameters / stuff that does not require gradient’<br>
parameters = ‘learnable parameters, requires gradient’</p>

111201: <p>Sort of hijacking the thread, but I am struggling at implementing capsule net, there is a need for some non-trainable variables, and unwanted in case of <code>state_dict</code>. Since those are just computed statistics.<br>
So the problem is those variables are in the model code, which I use code like <code>torch.zeros(b, h, w).cuda()</code>.<br>
But this is ugly, and if use ‘torch.zeros(b, h ,w)’, these variables will not be sent to GPU as we do <code>model.to(device)</code>.<br>
Please let me know if there is a better way to construct them.<img src="https://discuss.pytorch.org/images/emoji/apple/grimacing.png?v=9" title=":grimacing:" class="emoji" alt=":grimacing:"></p>

ptrblck: <p>Could you describe the usage of these tensors a bit?<br>
I assume they are not defining the model state, as you don’t want to have them in the <code>state_dict</code>, which means these tensors are independent of the model?<br>
Could you create these tensors then during runtime, e.g. by using the device attribute of a parameter or buffer?</p>

111201: <p>Do you mean something like</p>
<pre><code class="lang-auto">independent_tensor = torch.zeros(3, 3).to(feat_map.device)
</code></pre>
<p>yeah, it’s a nicer workaround. Thanks.<br>
But it will be better if there is a way to do this without setting device in the model part of code. So the whole model can be send to GPU or CPU as we set <code>model.to(device)</code></p>

ptrblck: <p><code>model.to()</code> transfers all “states” to the specified device.<br>
However, your use case seems as if the mentioned tensors should not be in the <code>state_dict</code>, which seems like a special use case.<br>
Could you therefore explain the use case a bit, i.e.:</p>
<ul>
<li>what are these tensors used for</li>
<li>are they specific to the model</li>
<li>how do you create them (model dependent or not?)</li>
</ul>

111201: <p>sorry about the vagueness.<br>
In the example of capsule net:</p>
<p><a href="https://github.com/gram-ai/capsule-networks/blob/master/capsule_network.py#L67" rel="noopener nofollow ugc">a example of these tensors in capsule net implementation</a></p>
<ul>
<li>
<p>there tensors are used for computing coefficients assigning to feature maps which produce these coefficients under <code>torch.no_grad()</code>.</p>
</li>
<li>
<p>these tensors have nothing to do with training or learning. Just some values computed by a certain procedure (dynamic routing in capsule net).</p>
</li>
<li>
<p>can simply seen as computing the cosine similarity of certain layers’ feature maps.</p>
</li>
</ul>
<p>And as you mentioned:</p>
<blockquote>
<p><code>model.to()</code>  transfers all “states” to the specified device.</p>
</blockquote>
<p>It seems <code>model.to()</code> only cares about the “states”. Maybe <code>some_tensor.to(feature_map.device)</code> is the best we can get.<img src="https://discuss.pytorch.org/images/emoji/apple/grimacing.png?v=12" title=":grimacing:" class="emoji" alt=":grimacing:" loading="lazy" width="20" height="20"></p>

jtencer: <p>You could overwrite the to or apply methods for your module to include transferring that specific tensor.  This way you would not have to pass the device to any additional parts of your module.</p>

yclin: <p>Hi, one more question:<br>
I have a huge tensor (700MB, precomputed, requires_grad=False) which is used for  tensor multiplication somewhere as a Module (as shown in the snippet)</p>
<p>When training the model with multiple GPUs, I need to push it to all GPUs. The easiest way would be using regist_buffer in a module. However this means the stat_dict would be larger than 700MB ( definitely not a good idea). So I was wondering the best way to push such a large tensor to all GPUs?</p>
<p>BTW, if I simply use “tensor.to(device)” , is the tensor gonna be pushed to all GPUs or only the default one?  (Had a test, seems like it is on the defalult gpu not all gpus.)</p>
<p>Thanks in advance!</p>
<pre><code class="lang-auto">class NewModule(nn.Module):

    def __init__(self, pre_matrix):
        super(NewModule, self).__init__()
        # Pre_matrix: NXP, of size 700MB, requires_grad=False
        self.pre_matrix = pre_matrix
        self.pre_matrix.requires_grad=False
        # self.register_buffer('pre_matrix', pre_matrix)  ### this means the stat_dic is larger than 700MB

    def forward(self, input):
        # input: MXN, on multiple gpus
        # output: MXP, on multiple gpus
        out = input @ self.Pre_matrix
        return out
</code></pre>

ptrblck: <p>You could still use <code>register_buffer</code> and set <code>persistent</code> to <code>False</code>, which won’t add this buffer to the <code>state_dict</code> as described in the <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_buffer">docs</a>.</p>