How to implement a deep ensemble

[cemde]

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?

[carmocca]

I see two potential options.

1. 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

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

[cemde]

Thank you. 1) is what I coded now.

[tchaton]

You could create 1 optimizer with n param groups too.

[cemde]

@tchaton great idea. I wrote

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

But what do I return from training_step?
A List of my n losses raises the following error In automatic optimization, `training_step` must either return a Tensor, a dict with key 'loss' or None.

Do I need to use manual optimization?

[cemde]

Here a solution with caching predictions:

class DeepEnsemble(pl.LightningModule):
    def __init__(self, cfg):
        super().__init__(cfg)
        self.net = nn.ModuleList([configure_network(self.cfg) for _ in range(self.cfg.METHOD.ENSEMBLE)])
        self.cache_preds = []
        
    def configure_optimizers(self):
        return [torch.optim.Adam(net.parameters(), lr=self.cfg.SOLVER.LR) for net in self.net]

    def forward(self, x, idx = None):
        if idx is None:
            x = torch.stack([net.forward(x) for net in self.net], dim=-1)
        else:
            x = self.net[idx].forward(x)
        return x

    def training_step(self, batch, batch_idx, optimizer_idx):
        image, label = batch["image"], batch["label"]
        logits = self.forward(image, optimizer_idx)
        self.cache_preds.append(logits.detach())

        loss = self.criterion(logits, label)
        
        if optimizer_idx == self.cfg.METHOD.ENSEMBLE - 1:
            logits = torch.stack(self.cache_preds, dim=-1)
            mean_logit = logits.mean(dim=-1)
            all_loss = self.log_loss(mean_logit, label, 'train')
            metrics = self.log_metrics(mean_logit, label, 'train')
            
            self.cache_preds, self.cache_loss = [], []

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

        mean_logit = logits.mean(dim=-1)

        loss = self.log_loss(mean_logit, label, 'val')
        metrics = self.log_metrics(mean_logit, label, 'val')
        return metrics[self.cfg.CKPT.MONITOR]