DataModule conditional dataloader not being able to recognize

[chris-tkinter]

Hi, I have a datamodule that looks like this

 def load_dataloader(self, phase: str):
        # when no test dataset
        if phase == "test" and not self.use_test:
            return None
        # self.datasets str -> torch.data.Dataset
        dataset = self.datasets[phase]
       # in meta learning there is no dataloader
       # but must sample several times
        if self.train_strategy == "meta_learning":
            for _ in range(self.meta_batch_size):
                yield dataset.getSamples()
        # baseline, non-meta-learning
        return DataLoader(
            dataset, batch_size=self.batch_size,
            num_workers=self.num_workers,
            pin_memory=True
        )

    def train_dataloader(self):
        return self.load_dataloader("train")

    def val_dataloader(self):
        print("val", self.batch_size, self.train_strategy)
        return self.load_dataloader("valid")

    def test_dataloader(self):
        return self.load_dataloader("test")

However, I found that withtrain_strategy == "meta_learning" or not, it would not go through the function load_dataloader, in fact, in here it would make the variable dataloader to be a python generator. As the result, the trainer somehow terminates directly without actual training.

I am wondering is there anything I did wrong, and what would be the best way to make a conditional dataloader?

A little more context: I am trying to implement a meta learning approach, and would like to compare the meta learning approach with the vanilla / traditional approach, therefore I have a condition train_strategy here.

Thanks!

[tchaton]

Hey @chris-tkinter,

Lightning expects for you to return a DataLoader or Iterator.

I recently implemented a MetaLearningDataset for distributed training: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/tests/utilities/test_auto_restart.py#L354.

I am working with the learn2learn team to provide better support for MetaLearning with Lightning.

Mind giving it a try ?

Best,
T.C