Package: torchmtl
A lightweight module for Multi-Task Learning in pytorch.torchmtltries to help you composing modular multi-task architectures with minimal effort. All you need is a list of dictionaries in which you define your layers and how they build on each other. From this,torchmtlconstructs a meta-computation graph which is executed in each forward pass of the createdMTLModel. To combine outputs from multiple layers, simplewrapper functionsare provided.Installationtorchmtlcan be installed viapip:pip install torchmtlQuickstart (or find exampleshere)Assume you want to train a network on three tasks as shown below.To construct such an architecture withtorchmtl, you simply have to define the following listtasks=[{'name':"Embed1",'layers':Sequential(*[Linear(16,32),Linear(32,8)]),# No anchor_layer means this layer receives input directly},{'name':"Embed2",'layers':Sequential(*[Linear(16,32),Linear(32,8)]),# No anchor_layer means this layer receives input directly},{'name':"CatTask",'layers':Concat(dim=1),'loss_weight':1.0,'anchor_layer':['Embed1','Embed2']},{'name':"Task1",'layers':Sequential(*[Linear(8,32),Linear(32,1)]),'loss':MSELoss(),'loss_weight':1.0,'anchor_layer':'Embed1'},{'name':"Task2",'layers':Sequential(*[Linear(8,64),Linear(64,1)]),'loss':BCEWithLogitsLoss(),'loss_weight':1.0,'anchor_layer':'Embed2'},{'name':"FNN",'layers':Sequential(*[Linear(16,32),Linear(32,32)]),'anchor_layer':'CatTask'},{'name':"Task3",'layers':Sequential(*[Linear(32,16),Linear(16,1)]),'anchor_layer':'FNN','loss':MSELoss(),'loss_weight':'auto','loss_init_val':1.0}]You can build your final model with the following lines in which you specify from which layers you would like to receive the output.fromtorchmtlimportMTLModelmodel=MTLModel(tasks,output_tasks=['Task1','Task2','Task3'])This constructs ameta-computation graphwhich is executed in each forward pass of yourmodel. You can verify whether the graph was properly built by plotting it using thenetworkxlibrary:importnetworkxasnxpos=nx.planar_layout(model.g)nx.draw(model.g,pos,font_size=14,node_color="y",node_size=450,with_labels=True)The training loopYou can now enter the typicalpytorchtraining loop and you will have access to everything you need to update your model:forX,yindata_loader:optimizer.zero_grad()# Our model will return a list of predictions (from the layers specified in `output_tasks`),# loss functions, and regularization parameters (as defined in the tasks variable)y_hat,l_funcs,l_weights=model(X)loss=0# We can now iterate over the tasks and accumulate the lossesforiinrange(len(y_hat)):loss+=l_weights[i]*l_funcs[i](y_hat[i],y[i])loss.backward()optimizer.step()Details on the layer definitionThere are 6 keys that can be specified (nameandlayersmustalways be present):layersBasically takes anynn.Modulethat you can think of. You can plug in a transformer or just a handful of fully connected layers.anchor_layerThis defines from which other layer this layer receives its input. Take care that the respective dimensions match.lossThe loss function you want to compute on the output of this layer (l_funcs). Can be set toNoneor omitted altogether when only access to the layer's output is needed.loss_weightThe scalar with which you want to regularize the respective loss (l_weights). If set to'auto', ann.Parameteris returned which will be updated through backpropagation. Can be set toNoneor omitted altogether when only access to the layer's output is needed.loss_init_valOnly needed ifloss_weight: 'auto'. The initialization value of theloss_weightparameter.Wrapping functionsNodes of themeta-computation graphdon't have to be pytorch Modules. They can beconcatenationfunctions or indexing functions that return a certain element of the input. If yourXconsists of two types of input dataX=[X_1, X_2], you can use theSimpleSelectlayer to select theX_1by settingfromtorchmtl.wrapping_layersimportSimpleSelect{...,'layers'=SimpleSelect(selection_axis=0),...}It should be trivial to write your own wrapping layers, but I try to provide useful ones with this library. If you have any layers in mind but no time to implement them, feel free toopen an issue.Cite@misc{torchMTL: A lightweight module for Multi-Task Learning in pytorch,
  author = {Bock, Christian},
  doi = {10.5281/zenodo.4362515},
  url = {https://github.com/chrisby/torchMTL},
  year = {2020}
}CreditsLogo credits and license: I reused and remixed (moved the dot and rotated the resulting logo a couple times) the pytorch logo fromhere(accessed throughwikimedia commons) which can be used under theAttribution-ShareAlike 4.0 Internationallicense. Hence, this logo falls under the same license.