Instantiate data augmentations through CLI
Hi,
I want to implement a BYOL-like model using Lightning CLI, and one of the key aspects is to use specific data augmentations. However, I don't manage to indicate to the cli that I want to use these specific data augmentations. Here is a snippet of my code:
main.py
import torch
import torch.nn as nn

import pytorch_lightning as pl
from pytorch_lightning.utilities.cli import LightningCLI, MODEL_REGISTRY
from pl_bolts.datamodules import CIFAR10DataModule
from pl_bolts.models.self_supervised.simclr import SimCLRTrainDataTransform, SimCLREvalDataTransform


class DummyModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(32*32*3, 10)
        self.loss_fn = nn.MSELoss()

    def shared_step(self, batch, batch_idx):
        x, y = batch[0][:2]
        z1 = self.fc(x.reshape(x.size(0), -1))
        z2 = self.fc(y.reshape(y.size(0), -1))
        return self.loss_fn(z1, z2)

    def training_step(self, batch, batch_idx):
        return self.shared_step(batch, batch_idx)

    def validation_step(self, batch, batch_idx):
        return self.shared_step(batch, batch_idx)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters())


# THE COMMENTED LINES BELOW RUN PERFECTLY
# trainer = pl.Trainer()
# model = DummyModel()
# dm = CIFAR10DataModule()
# dm.train_transforms = SimCLRTrainDataTransform(32)
# dm.val_transforms = SimCLREvalDataTransform(32)
# trainer.fit(model, dm)

cli = LightningCLI(DummyModel, CIFAR10DataModule, run=False)

# not instantiated!
print(cli.config_init.data.train_transforms)
and here is my config.yaml:
data:
  train_transforms:
    class_path: pl_bolts.models.self_supervised.simclr.SimCLRTrainDataTransform
    init_args:
      input_height: 32
  val_transforms:
    class_path: pl_bolts.models.self_supervised.simclr.SimCLREvalDataTransform
    init_args:
      input_height: 32
To run the code, I run the following command:
python main.py --config config.yaml

and here is the error I get:
Traceback (most recent call last):
  File "/home/alain/code/misc/examples/main.py", line 42, in <module>
    cli.trainer.fit(cli.model, cli.datamodule)
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 740, in fit
    self._call_and_handle_interrupt(
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 685, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 777, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1199, in _run
    self._dispatch()
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1279, in _dispatch
    self.training_type_plugin.start_training(self)
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 202, in start_training
    self._results = trainer.run_stage()
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1289, in run_stage
    return self._run_train()
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1311, in _run_train
    self._run_sanity_check(self.lightning_module)
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1375, in _run_sanity_check
    self._evaluation_loop.run()
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/pytorch_lightning/loops/base.py", line 145, in run
    self.advance(*args, **kwargs)
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py", line 110, in advance
    dl_outputs = self.epoch_loop.run(dataloader, dataloader_idx, dl_max_batches, self.num_dataloaders)
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/pytorch_lightning/loops/base.py", line 140, in run
    self.on_run_start(*args, **kwargs)
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 86, in on_run_start
    self._dataloader_iter = _update_dataloader_iter(data_fetcher, self.batch_progress.current.ready)
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/pytorch_lightning/loops/utilities.py", line 121, in _update_dataloader_iter
    dataloader_iter = enumerate(data_fetcher, batch_idx)
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/pytorch_lightning/utilities/fetching.py", line 199, in __iter__
    self.prefetching(self.prefetch_batches)
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/pytorch_lightning/utilities/fetching.py", line 258, in prefetching
    self._fetch_next_batch()
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/pytorch_lightning/utilities/fetching.py", line 300, in _fetch_next_batch
    batch = next(self.dataloader_iter)
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 521, in __next__
    data = self._next_data()
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 561, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 49, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 49, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/torch/utils/data/dataset.py", line 363, in __getitem__
    return self.dataset[self.indices[idx]]
  File "/home/alain/miniconda3/envs/ts/lib/python3.9/site-packages/torchvision/datasets/cifar.py", line 121, in __getitem__
    img = self.transform(img)
TypeError: 'dict' object is not callable

What I understand from the error is that LightningCLI doesn't manage to interprete and instantiate the data augmentation model and understands it only as a dictionary with keys class_path and init_args. Does someone encounter a similar issue and/or knows how to solve it?
cc @carmocca