question: How to design a multi-task architecture?
context: 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 😨
answer: 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