Correct way to set up multi-task training?

[turian]

We have a NeuroIPS competition in submission, and I would like to use pytorch-lightning for the dev-kit.

I'm trying to figure out the most elegant pattern for using ptl to learn an embedding in a multi-task learning scenario:

1. There is an embedding module that takes the input and projects it to a low dimensional representation e(x).
2. There are K different learning tasks, some of which have MANY examples and some of which have FEW examples. Note that the input has a different tensor size for each task, so we can't combine dataloaders. Consider, e.g., one task uses 4-second audio, one task uses 1-second audio, etc.
3. Each task uses e(x) as the input and adds a few layers, fine-tuning e(x) in the process.
4. We would prefer to cycle between the tasks evenly, rather than letting large tasks run for one epoch and then switching tasks.

Is there an elegant pattern for interleaving mini-batching training of multiple tasks simultaneously, all of which use a shared embedding module?