问如何进行多头学习

EN

#please note this is just pseudocode and I'm not well versed with computer vision
#therefore you need to check if resnet50 import is correct and look 
#for the imports of the task specific stuff
from torch import nn
from torchvision.models import resnet50

class MultiTaskModel(nn.Module):
    def __init__(self):
        #shared part
        self.resnet50 = resnet50()

        #task specific stuff
        self.tasks = nn.ModuleList()
        self.tasks.add_module('depth', Depth())
        self.tasks.add_module('denseflow', Denseflow())
        #...

    def forward(self, tasktag, ...):
        #shared part
        resnet_output = self.resnet50(...)

        #task specific parts
        if tasktag == 'depth':
            return self.tasks.depth(resnet_output)
        elif tasktag == 'denseflow':
            return self.tasks.denseflow(resnet_output)
        #...