Multiple models, one dataloader?

[turian]

I have a training regime that is disk-speed bound, because instances are loaded from disk.

I would like to train multiple models with one dataloader. That way, I can do model selection over many models, but reduce the number of disk reads. Is this possible?

[tchaton]

Dear @turian,

Yes, it is possible.

You could do something like this.

class MultiModels(LightningModule):

    def __init__(self, models: List[nn.Module]):
        self.models = models

    def compute_loss(self, model, batch):
        loss = ...
        return loss

    def training_step(self, batch, batch_idx):
        loss = sum(compute_loss(model, batch) for model in self.models)
        return loss


model = MultiModels([resnet50_model, alexnet_model, ...])
dm = ...
trainer.fit(model, dm)

Does this answer your questions ?