Defining modules in init but not using in forward method

KUSHAL_JAIN: <p>I have a neural net wherein I define multiple nn.Sequential modules in the <code>__init__()</code>. However, I use only one of them in the forward method based on a condition. Still, for all these modules, the weights and biases are initialized and <code>requires_grad</code> is also <code>True</code> which is expected. I do not actually need the gradients of modules that I do not use. Example code:</p>
<pre><code class="lang-auto">class Net(nn.Module):
  def __init__(self, head_type):
        self.head_type = head_type
        self.linear_head = nn.Linear(in_features=model_dim, out_features=classifier_dim)
        
        self.lstm = nn.LSTM(model_dim, model_dim, batch_first=True)
        
        self.multilinear_head = nn.Sequential(nn.Linear(in_features=model_dim, out_features=256),
                                               nn.Dropout(0.2),
                                               nn.Linear(in_features=256, out_features=classifier_dim),
                                               nn.Dropout(0.2)
                                              )

 def forward(self, x):
      if self.head_type == 'linear':
         out = self.linear_head(x)
     return x 
</code></pre>
<p>Is this alright with respect to space the additional tensors are taking? I guess since forward pass is not peformed for the unused modules they would be taking minimal space. Any work-around for this would be helpful. Thank you</p>

ptrblck: <p>Since some modules are not used during the forward pass, no intermediate tensors would be created. However, the parameters of these unused modules will still be pushed to the device.<br>
If you don’t want to use these modules, you could either remove them from the <code>__init__</code> or delete them afterwards:</p>
<pre><code class="lang-python">model.multilinear_head = None
</code></pre>
<p>The cleaner solution would be the first approach, as the second one might have side effects (e.g. I don’t know, if forward hooks would break etc.).</p>
<p>Alternatively, you could also manually move the unused modules back to the CPU (which would at least save some GPU memory):</p>
<pre><code class="lang-python">model.multilinear_head.cpu()
</code></pre>

KUSHAL_JAIN: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="92442">
<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 don’t want to use these modules, you could either remove them from the <code>__init__</code></p>
</blockquote>
</aside>
<p>If I remove them from <code>__init__</code>, where do I define them? Since what head has to be chosen depends on the conditional parameter that will be passed to the <code>__init__</code> i.e. the <code>head_type</code>. What I could perhaps do is make separate methods within the class for each head and put conditions in the <code>__init__</code> itself. Any other cleaner solution for this?</p>

ptrblck: <p>I think a conditional creation of the desired “head” module in <code>__init__</code> sounds like a proper way.</p>