Pattern for specifying modules in argument parser versus CLI

[laserkelvin]

With the new experimental CLI, I've been trying to think of how to reconcile using the add_model_specific_args pattern with YAML configs; I've been digging through the docs and couldn't find the pieces I needed so hoping others could chime in. Basically I'd want to implement this once, and have backwards compatibility for both the argparse and new CLI routes, and I might be mistaken, but it feels that the two routes involve mutually exclusive patterns.

Let's say we have an abstract pl.LightningModule that allows the user to choose an activation function:

from ast import literal_eval
from torch import nn
import pytorch_lightning as pl

class LitModule(pl.LightningModule):
    def __init__(self, input_dim: int, output_dim: int, activation: str):
        # cross your fingers it resolves, something like `nn.SilU`
        act_class = literal_eval(activation)
        # instantiate the activation function object
        self.model = nn.Sequential(nn.Linear(input_dim, output_dim), act_class())

With argparse, the way I would set this up would be to implement add_model_specific_args:

...
    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = parent_parser.add_argument_group("LitModule")
        parser.add_argument("--input_dim", type=int)
        parser.add_argument("--output_dim", type=int)
        parser.add_argument("--activation", type=str, default="nn.SiLU", help="Reference to activation function class in `torch.nn`.")

And in my training script:

from torch import nn

parser = ArgumentParser()
parser = LitModule.add_model_specific_args(parser)

args = parser.parse_args()

model = LitModule(args)

With the new CLI, the better option would be to use nn.Module for typing instead of str for the activation argument, i.e. redefining the __init__ so we remove the possibility of running arbitrary code via literal_eval since it's taken care of by jsonargparse/LightningCLI:

class LitModule(pl.LightningModule):
    def __init__(self, input_dim: int, output_dim: int, activation: nn.Module):
        # don't need to use `literal_eval` anymore
        self.model = nn.Sequential(nn.Linear(input_dim, output_dim), activation())

...and in my YAML config:

model:
   class_path: mymodule.LitModule
   init_args:
      input_dim: 8
      output_dim: 2
      activation: torch.nn.SiLU

I can't think of any way to implement this same pattern with argparse. While I prefer the CLI approach because of its potential for composability, it's experimental and sometimes you want more granular control as you get from a training script, i.e. putting the components together yourself, multiple optimizers, etc. which has a bit more flexibility. There might not be a one-size fits all, but I'd appreciate any feedback/discussion.

[mauvilsa]

If you truly want to have both, you could do activation: Union[str, nn.Module] and use the literal_eval only if an str is received. But even if you do this, there is little reason to keep using add_model_specific_args since with LightningCLI you could also give an str instead of the class_path and init_args pair.

Regarding "i.e. putting the components together yourself, multiple optimizers, etc. which has a bit more flexibility" best if you give more details about what you want to do, and then see if it makes sense or not to use LightningCLI.

[laserkelvin]

Thanks a lot for your insights!

I don't think I had concrete examples of what I need to achieve immediately with my comment about multiple optimizers (I guess GANs come to mind most readily), but more broadly speaking wondering whether or not the two ways of using Lightning (argparse versus LightningCLI) were mutually exclusive, and for the sake of future proofing models we write now, whether we should design them for one way or the other.

Just to go back on optimizers, the older (argparse) route was to define lr/weight_decay as optional arguments:

class LitModule(pl.LightningModule):
    def __init__(self, input_dim: int, output_dim: int, activation: str, lr: Optional[float] = 1e-3, weight_decay: Optional[float] = 0.):
        ...
        self.save_hyperparameters()

    def configure_optimizers(self):
        optimizer = optim.Adam(self.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay)
        # optionally define a scheduler, etc.

whereas in the CLI configuration, we don't have to define configure_optimizers at all, and the optimizer arguments are no longer needed for LitModule. One could keep the definition, but then your code isn't as clean as it could be, and it could lead to unexpected behavior (e.g. if you leave an optimizer definition in the YAML by accident).

At this point I think this is probably largely a matter of stylistic choice, but I'd appreciate any feedback/discussion; this kind of thing isn't captured in the style guide yet.

[mauvilsa]

I would say to forget about the argparse route. It is very limiting and requires a bunch of boilerplate.

The automatic configure_optimizers is only to avoid the boilerplate for the simple cases in which this is always the same. But this is only optional. You can implement configure_optimizers in your module and make the optimizer(s) lr scheduler(s) configurable. Look at the documentation where it says "The automatic implementation of configure_optimizers can be disabled ..."