Track non-learnable parameters in PyTorch

jmaronas: <p>Hello.</p>
<p>I have a model in which some parameters are not learnable (e.g a PCA projection). I would like to know if there is a way to to track non-learnable parameters in PyTorch <code>nn.Module</code> so that when I execute <code>my_model.to(device)</code> all these parameters are moved between devices.</p>
<p>Why I need this behavior? Well my models are trained on GPU but evaluated on CPU for safeness, so that I dont run out of memory. This is because during evaluation I am performing a Monte Carlo expectation hence the predictive distributions use more samples during test than during training.</p>
<p>Thank you</p>

jmaronas: <p>My solution has been to add this method:</p>
<pre><code>def _custom_to(self,device):
    self.to(cg.device)
    for mod in self.modules():
        try:
            mod.__move_to_device__(device)
        except:
            pass
</code></pre>
<p>where the Module class has now a <code>mod.__move_to_device__</code> that moves to the desired device the desired variables (in my case the non learnable parameter). I am not sure if there is a better way to do this. Also not sure if it worth working on incorporating this feature to the library. It would be as easy as adding the <code>_custom_to()</code> method and the <code>__move_to__device</code>. The later calls <code>pass</code> by default and can be overwritten by the user as I have done. The other solution is to call <code>__move_to__device</code> directly inside the standard <code>nn.Module.to()</code> method (which I think is better)</p>

albanD: <p>Hi,</p>
<p>This is what buffers are for.<br>
You can do <code>self.register_buffer("foo", your_tensor)</code>. You can then use it as if you did <code>self.foo = your_tensor</code>.<br>
And it will be moved around for you and saved into the state dict as well.</p>

jmaronas: <p>Thanks.</p>
<p>But if saved into the state dict then it means I have to remove it when passing nn.Parameters() into an optimizer right? As far as I remember the optimizer raise error for tensors that do not require grad.</p>

albanD: <p>You pass <code>mod.parameters()</code> to the optimizer right? While the state dict is <code>mod.state_dict()</code> which is what you save. And what you save contains the buffers.</p>
<p>Usual buffers we use are for example the running statistics of batchnorm.</p>

jmaronas: <p>Yes, well in my case I don’t really need to save this parameter into my state_dict. I just need that this parameter is moved between devices when call <code>model.to(device)</code> and that is not passed to the optimizer when calling <code>model.parameters()</code>.</p>
<p>I could workaround this issue if I run in the forward of my nn.Module something like:</p>
<pre><code class="lang-auto">W = self.W.to(self.device)
Y = X@W # here X can be in GPU or CPU depending on the iteration 
</code></pre>
<p>But I would prefer not calling the method <code>to.(self.device)</code> on each forward and just handle the device of each of the parameters through a call to <code>my_model.to()</code></p>

albanD: <p>Again, buffers are <em>never</em> in the <code>model.parameters()</code> !! only parameters are there.</p>
<p>But if you don’t want a buffer to be part of the state dict (which is different from .parameters()!) you can just give <code>persistent=False</code> when you create it <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"><br>
You can check the <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_buffer">doc</a> for more details.</p>

jmaronas: <p>Thank you. Sorry I was missunderstanding. Now is clear</p>