How to design a multi-task architecture?

[teufelweich]

Hello,

I'm designing a multi-task architecture for depth estimation and semantic segmentation. Because these are very similar tasks, I can use my existing nn.Module to learn both of them. I wrapped the nn.Module in a pl.LightningModule and got a working segmentation. Now I want to extend the LightningModule to be able to do the depth estimation.
I want to reuse as much code as possible to reduce errors. But the validation for both tasks is very different, because I'm using many metrics and extensive W&B logging. I first thought of subclassing the pl.LightningModule to a segmentation and a depth estimation module, but I want to be able to use both segmentation and depth in a single LightningModule later on, when the network should learn both tasks.
Thus I created methods for the specific validation steps and set them to variables in the classes init. These get called from the overwritten methods of the class.
Heres some pseudo-code from my structure:

class LitModel(pl.LightningModel):
 def __init__(self, model: nn.Module,  hparams):
        super().__init__()
        # default for segmentation and depth
        self.hparams.update(hparams)
        self.model = model

        # SEGMENTATION STUFF
        if segmentation:
            
            self._batch_target_name = 'label'

            # VAL STUFF
            self._loss_func = CrossEntropyLoss
            self._val_loss_func = CrossEntropyLoss
            self._confusion_matrix = ConfusionMatrix(num_classes=num_classes, compute_on_step=False)
            self._best_metric = 0
            self._best_metric_name = 'mIoU'

            self._val_reset_metric_list = [self._val_loss_func, self._confusion_matrix]

            self._validation_step = self.seg_validation_step
            self._validation_epoch_end = self.seg_validation_epoch_end

        # DEPTH STUFF
        elif depth:
            
            self._batch_target_name = 'depth'

            # VAL STUFF
            self._loss_func = L1
            self._val_loss_func = MSE
           
            self._best_metric = 0
            self._best_metric_name = 'RMSE'

            self._val_reset_metric_list = [self._val_loss_func]

            self._validation_step = self.dpt_validation_step
            self._validation_epoch_end = self.dpt_validation_epoch_end

    def forward(self, sample):
        return self.model(sample)

    def training_step(self, batch, batch_idx):
        # disassemble batch/samples
        image = batch['image']
        target_scales = [batch[self._batch_target_name]]

        # predict input images
        pred_scales = self.model(image)

        # calculate losses
        losses = self._loss_func(pred_scales, target_scales)
        summed_loss = sum(losses)
        self.log('train/loss', summed_loss, on_step=False, on_epoch=True)  # logs mean of all losses during that epoch
        return {'loss': summed_loss}

    def validation_step(self, batch, batch_idx):
        # disassemble batch/samples
        image = batch['image']
        gt = batch[self._batch_target_name]

        # predict input images
        prediction = self.model(image)

        return self._validation_step(batch, batch_idx, gt, prediction)

    def seg_validation_step(self, batch, batch_idx, gt, prediction):
        # calculate segmentation validation losses

    def dpt_validation_step(self, batch, batch_idx, gt, prediction):...
        # calculate depth validation losses

    def on_validation_epoch_start(self) -> None:
        for f in self._val_reset_metric_list:
            f.reset()

    def validation_epoch_end(self, outputs) -> None:
        metric = self._validation_epoch_end(outputs)
        if self._best_metric < metric:
            self.log(f'eval/best_{self._best_metric_name}', metric)
            self.log('eval/best_epoch', self.current_epoch)
            self._best_metric = metric

    def seg_validation_epoch_end(self, outputs) -> metric:
        # calculate confusion matrix and mIoU
        return mIoU

    def dpt_validation_epoch_end(self, outputs) -> metric:
        # calculate RMSE
        return RMSE

    def configure_optimizers(self):...

To me, this seems as really bad practice and I hope it doesn't hurt you to much seeing this.
What would be a better/the best way to achieve this? I'm looking for the best practice.

I'm coming from TensorFlow and I really really like how clean PyTorch Lightning is. It feels so good to work with it. But what I wrote up there just feels wrong, but it works 😨

[justusschock]

First of all: I'm glad you are enjoying lightning :)

Coming to your code: It actually doesn't look so bad to me. You're separating functionality for different use cases in different functions (which is perfectly fine). What you could do ( if you really want to) is something like this:

class BaseModel(LightningModule):
   ... # implements all the logic to be shared between the models such as the module logic or something like this

class SegmentationModel(BaseModel):
    ... # adds all the segmentation-only logic

class DepthModel(BaseModel):
    ... # adds all the depth-only logic

class CombinedModel(LightningModel):
    def __init__(self, model, hparams):
        if depth:
            model = DepthModel(model, hparams)
        else:
            model = SegmentationModel(model, hparams)
        self.model = model

    def training_step(self, *args, **kwargs):
        return self.model.training_step(*args, **kwargs)

    # do the same for other methods and hooks

That way your classes would be a bit more separated and self-contained. That being said, I still think your current approach is perfectly fine

[teufelweich]

Thanks for your reply.

This is what I thought of initially. But later on, I want to predict the segmentation and depth simultaneously, sharing the encoder and having separated decoders. Thus I would call the specific depth and segmentation functions inside the CombinedModels validation step. This would not work if the logic is directly in the overwritten methods inside the DepthModel and SegmentationModel. I just hoped there could be a way to combine the hooks/overwritten methods of two LightningModules. But this is probably a general abstraction of OOP that I'm missing here.

I'm thinking of a meta-architecture that allows arbitrary number of tasks, where I can define the tasks specifics for each hook in an elegant way. But that's really hard and I don't have the experience to come up with such abstractions. Probably should stick to a combination of subclassing and the specific methods because it works.

[justusschock]

@teufelweich I think there is no general way to do so. That probably depends a lot on your data format and how you load it.

Do you have one loader providing one image and targets for both, depth estimation and segmentation?

[teufelweich]

Yes, the DataModule provides everything in each sample.

[justusschock]

I think then using your approach is probably the easiest (and cleanest in terms of less boilerplate)

[teufelweich]

Great. Thanks for your assessment. :)