How to implement a deep ensemble
I am looking to implement n parallel independent ensembles. My idea is the following:
class DeepEnsemble(LightningModule):
    def __init__(self, cfg):
        super().__init__(cfg)
        self.net = nn.ModuleList([configure_network(self.cfg) for _ in range(self.cfg.METHOD.ENSEMBLE)])

    def configure_optimizers(self):
        return [torch.optim.Adam(net.parameters(), lr=self.cfg.SOLVER.LR) for net in self.net]

    def forward(self, x):
        x = [net.forward(x) for net in self.net]
        return x

    def training_step(self, batch, batch_idx, optimizer_idx):
        image, label = batch["image"], batch["label"]
        logits = self.forward(image)
        loss = [self.criterion(logit, label) for logit in logits]

        mean_logit = torch.stack(logits, dim=-1).mean(dim=-1)

        metrics = self.log_metrics(mean_logit, label, 'train')

        return loss

    def validation_step(self, batch, batch_idx):
        image, label = batch["image"], batch["label"]
        logits = self.forward(image)

        mean_logit = torch.stack(logits, dim=-1).mean(dim=-1)

        metrics = self.log_metrics(mean_logit, label, 'val')

        return metrics[self.cfg.CKPT.MONITOR]

    def test_step(self, batch, batch_idx):
        pass
I have n networks and n optimisers. My solution works (I think), but the training_step gets called with a new optimizer_idx every time, which indicates that Pytorch Lightning expects to only train 1 network per training_step. Therefore, my solution is very inefficient, because n^2 forward passes are executed instead of n. If I only do the forward pass for the ith network, then I can't compute metrics based on all ensembles (e.g. disagreement) unless I write some very inelegant if statements.
In addition It would be nice to have all forward passes done in parallel instead of sequential like in this list comprehension.
So what is the most elegant way to train an ensemble and still access all predictions for metric logging together?
I see two potential options.


cache the forward output for a specific batch idx. Check the automatic optimization flow: https://pytorch-lightning.readthedocs.io/en/latest/common/optimizers.html#automatic-optimization


Use manual optimization. https://pytorch-lightning.readthedocs.io/en/latest/common/optimizers.html#manual-optimization