[Retiarii] Usage of DataLoaders in one-shot NAS Trainers

[thomasschmied]

Describe the issue:

Hello!

I've been working with NNI recently and I really like the Retiarii Mutation API.

While running one-shot NAS experiments on custom datasets with ENAS and Darts, I encountered a problem. All Trainer classes in nni.retiarii.oneshot.pytorch, (e.g., ENASTrainer) construct the torch.utils.data.DataLoader instances in _init_dataloader. Furthermore, the _init_dataloader function does a 50:50 split of the given PyTorch Dataset instance to construct the train and validation sets.

However, this behaviour is a bit limiting. It is particularly problematic if custom datasets or datasets with predefined train-valid-test splits are used. Therefore, my question: would it be possible to change this behaviour?

Possible solutions:

• Pass train and valid DataLoader instances directly to the Trainer class instead of constructing them in _init_dataloader
• Remove _init_dataloader from the Trainer __init__ and make the function configurable for the user

Currently, I bypassed this problem by overwriting the behaviour of _init_dataloader.
However, I believe that these changes would make the library more generally applicable to a broader range of use cases. Not sure if other people encountered this problem before.

Environment:

• NNI version: 2.3
• Python version: 3.8.8
• PyTorch: 1.9.0

Thank you,
Thomas

[ultmaster]

We might refactor all the one-shot trainers from a high-level perspective, and make customization more friendly. We will consider your use case when we do that. Meanwhile, the recommended practice is to copy/inherit the existing one-shot trainers and adapt to your own use scenario.