How to Multi-Head learning

Question

I have about 5 models that work pretty well trained individually but I want to fuse them together in order to have one big model. I'm looking into it because one big model is more easy to update (in production) than many small model this is an image of what I want to achieve.

my question are, is it ok to do it like this ? having one dataset per head model, how am I supposed to train the whole model ?

Answer 1

my question are, is it ok to do it like this

Sure you can do that. This approach is called multi-task learning. Depending on your datasets and what you are trying to do, it will maybe even increase the performance. Microsoft used a multi-task model to achieve some good results for the NLP Glue benchmark, but they also noted that you can increase the performance further by finetuning the joint model for each individual task.

having one dataset per head model, how am I supposed to train the whole model?

All you need is pytorch ModuleList:

#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)
        #...