Training nnModuleList separately

Nerolf05: <p>I am trying to train a list of neural networks in parallel, using nn.ModuleList.</p>
<pre data-code-wrap="python"><code class="lang-python">class Modules(nn.Module):
    def __init__(self, n_models):
        super().__init__()
        self.modules = nn.ModuleList([nn.Sequential(nn.Linear(1, n), nn.Linear(n, 1)) for _ in range(n_models)])

    def forward(self, x):
        y = torch.cat([i_module(x) for i_module in self.modules], dim=-1)
        return y

...

def train_epoch(model, criterion, optimizer, train_dl):
    model.train()
    criterion.train()
    for data, target in train_dl:
        optimizer.zero_grad(set_to_none=True)        
        prediction = model(data)
        loss = criterion(input=prediction, target=target) # shape [n_modules x 1]
        loss = loss.sum() # loss.mean()
        loss.backward()
        optimizer.step() # update weights
    pass
...

</code></pre>
<p>During training, I would like to use the respective loss value to update the respective module only, and prevent <em>loss[0]</em> from affecting the update of <em>module[1]</em>.</p>
<p>Rather than reducing the loss vector to a scalar via <em>sum()</em> or <em>mean()</em> and affecting the optimization of all modules, I would like to use the actual loss value (<em>loss[i]</em>) to update the respective module (<em>self.modules[i]</em>) and obtain independently trained models.</p>
<p>I have read that <em>backward()</em>, if non-scalar, requires the gradients, which i am unsure how to compute manually, without messing up the computation graph.</p>
<p>Also, iterating the loss vector is probably not effective, since the full graph will be overwritten with each iteration.</p>
<pre data-code-wrap="python"><code class="lang-python">for i in range(loss.shape[0]):
    loss[i].backward(retain_graph=True)
optimizer.step()
</code></pre>
<p>The question is probably related to whether PyTorchs gradient</p>
<pre data-code-wrap="markdown"><code class="lang-markdown">$\frac{\delta loss.sum()}{\delta prediction[0]} == \frac{\delta loss[0]}{\delta prediction[0]}$
</code></pre>

Arunprakash-A: <p>If you treat a list of neural networks as independent, then each optimizer should also be independent. That is, you need to use a separate optimizer (e.g., <code>optimizer_i</code>) for each network</p>

Tanmay_Pani: <p>Your approach is fine but might get cumbersome for managing the optimizer states for a large amount of models. I would recommend setting different parameter groups in one optimizer. So you would define your optimizer as:</p>
<pre data-code-wrap="python"><code class="lang-python">optim = torch.optim.Adam(
            [{"params" : m.named_parameters()} for m in model.modules]
        )
</code></pre>
<p>And then in training step if you do:</p>
<pre data-code-wrap="python"><code class="lang-python">loss.backward()
optim.step()
</code></pre>
<p>Different modules in your model will be updated independently. You can also set different learning_rates for each parameter group. You can find more detail [here]( <a href="https://docs.pytorch.org/docs/stable/optim.html#taking-an-optimization-step" class="inline-onebox" rel="noopener nofollow ugc">torch.optim — PyTorch 2.8 documentation</a> ).</p>