I'm working on a Multi-Task Learning (MTL) setup and need to feed multiple datasets into my model. How does PyTorch Lightning handle multiple train dataloaders? Is there a specific way to aggregate the losses from different sources in the training_step to ensure gradients flow correctly?
3 answers
In PyTorch Lightning, you can return a dictionary or a list of DataLoaders in your train_dataloader() method. For MTL, the CombinedLoader class is particularly useful. In your training_step, the batch argument will mirror the structure of your loaders. If you passed a dictionary, batch will be a dictionary containing data from both tasks. You can then compute separate losses, weight them, and return the total loss. The Trainer handles the optimization step automatically. This approach ensures your PyTorch Lightning code stays organized even when the data pipeline gets complex, avoiding the messy manual loops typically found in vanilla PyTorch.
Does the CombinedLoader support different sampling strategies, like if one dataset is much smaller than the other? I'm worried about the smaller task being over-sampled.
Just use a dictionary in train_dataloader and access them by key in your training_step. It keeps the PyTorch Lightning logic very straightforward for task-specific losses.
Sandra's method is the most readable. Keeping task keys clear makes it easier for others to understand the MTL flow in your PyTorch Lightning implementation.
Gary, you can set the mode in CombinedLoader to "max_size_cycle" or "min_size". This tells PyTorch Lightning how to handle the discrepancy in dataset lengths, preventing the model from overfitting on the smaller dataset while waiting for the larger one to finish.