Customizing LightningCLI for multiple Test Stages

[Amerehei]

I'm working on customizing LightningCLI to include two separate test stages: "clean" and "robust." I'd appreciate guidance on how to achieve this effectively with the LightningCLI framework.

Consider the following code snippet:

import torchattacks
atk = torchattacks.PGD(model, eps=8/255, alpha=2/255, steps=4)
# If inputs were normalized, then
# atk.set_normalization_used(mean=[...], std=[...])
adv_images = atk(images, labels)

1. How can I set images = adv_images before running the validation?
   I want to modify the validation dataset so that adversarial examples are used during validation instead of the clean images. note that images and labels are a single batch

2. How can I switch between "clean" and "robust" modes from the CLI?
   I would like to be able to choose between clean validation and adversarial validation (robust mode) using command-line arguments.

3. Do I need to denormalize images?
   In other words, are inputs normalized in the line we modify?

I'm looking for the easiest and cleanest way to achieve these tasks while working within the LightningCLI framework.

[mauvilsa]

Just a recommendation from my side. LightningCLI just exposes parameters and type hints making them configurable through a CLI. Actual logic should not be implemented in the CLI. Better to think about what to implement in order to run things from python (without a CLI), but making sure that parameters are exposed and have good type hints. After this, it is likely that the CLI part will be almost automatic.

So my recommendation would be to figure out how to modify your LightningModule and LightningDataModule. Also you might want to subclass the Trainer class and modify the exposed Trainer methods.