[Retiarii] Rewrite trainer with PyTorch Lightning #3359
Conversation
dependencies/recommended.txt
Outdated
| @@ -1,6 +1,7 @@ | |||
| tensorflow | |||
| torch >= 1.6+cpu, != 1.7+cpu -f https://download.pytorch.org/whl/torch_stable.html | |||
| torchvision >= 0.8+cpu -f https://download.pytorch.org/whl/torch_stable.html | |||
| pytorch-lightning | |||
| onnx | |||
| peewee | |||
| thop | |||
There was a problem hiding this comment.
@colorjam do we still need thop?
There was a problem hiding this comment.
Will remove it in another PR.
| from torchvision import transforms | ||
|
|
||
| transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) | ||
| train_dataset = blackbox(MNIST, root='data/mnist', train=True, download=True, transform=transform) |
There was a problem hiding this comment.
blackbox is strange here, because i don't understand why it creates train_dataset
There was a problem hiding this comment.
o... MNIST is a class name?
There was a problem hiding this comment.
Yes. This is limited by serialization.
There was a problem hiding this comment.
is it possible we also wrap the dataset class by default? when users want to define their own dataset class, they decorate this class with for example @register_dataset
There was a problem hiding this comment.
another option, for this pr, is renaming blackbox to make_serializable
| def __init__(self, model, optimizer_class_name='SGD', learning_rate=0.1): | ||
| @blackbox_module | ||
| class AutoEncoder(LightningModule): | ||
| def __init__(self): |
There was a problem hiding this comment.
where is model configured?
| @register_trainer | ||
| class MnistTrainer(BaseTrainer): | ||
| def __init__(self, model, optimizer_class_name='SGD', learning_rate=0.1): | ||
| @blackbox_module |
There was a problem hiding this comment.
blackbox_module used on trainer is not that clear, let's discuss the name then
| applied_mutators: Mutator = None, strategy: BaseStrategy = None): | ||
| def __init__(self, base_model: nn.Module, trainer: Union[TrainingConfig, BaseOneShotTrainer], | ||
| applied_mutators: List[Mutator] = None, strategy: BaseStrategy = None): | ||
| # TODO: The current design of init interface of Retiarii experiment needs to be reviewed. |
There was a problem hiding this comment.
trainer could be both configuration and instantiated trainer?
There was a problem hiding this comment.
Have to admit that this is very confusing. Maybe we can find a better name than "TrainingConfig".
| @@ -1,3 +1,5 @@ | |||
| # This file is deprecated. | |||
There was a problem hiding this comment.
why not directly remove it?
There was a problem hiding this comment.
Because I don't want to remove the cgo_engine. If I remove this file and do not remove cgo_engine, linter will complain.
| Class for optimizer (not an instance). default: ``Adam`` | ||
| """ | ||
|
|
||
| def __init__(self, criterion: nn.Module = nn.MSELoss, |
There was a problem hiding this comment.
what is the difference between Regression and Classification?
|
the code looks great, one thing is that it is not easy to understand how to write/use trainer |
| class MnistTrainer(BaseTrainer): | ||
| def __init__(self, model, optimizer_class_name='SGD', learning_rate=0.1): | ||
| @blackbox_module | ||
| class AutoEncoder(LightningModule): |
There was a problem hiding this comment.
since we do not define model in LightningModule, if we still use the name LightningModule, it may be misleading. we can use more understandable name, for example, BaseTrainer, TrainingModule, etc.
| If the ``lightning_module`` has a predefined val_dataloaders method this will be skipped. | ||
| """ | ||
|
|
||
| def __init__(self, lightning_module: LightningModule, trainer: Trainer, |
There was a problem hiding this comment.
i am confused, what is the difference between lightning_module and trainer?
There was a problem hiding this comment.
let's think about how to support standalone mode, which would help us to think about how extensible is our current design
| With FunctionalTrainer | ||
| ^^^^^^^^^^^^^^^^^^^^^^ | ||
|
|
||
| There is another way to customize a new trainer with functional APIs, which provides more flexibility. Users only need to write a fit function that wraps everything. This function takes one positional arguments (model) and possible keyword arguments. In this way, users get everything under their control, but exposes less information to the framework and thus fewer opportunities for possible optimization. An example is as belows: |
There was a problem hiding this comment.
we should consistently use the word "trainer"
Work items
self-supervision).Future items