How to exclude model __init__ parameters from the lightning CLI?

[NiklasKappel]

I have a model that looks something like this:

class MyModel(L.LightningModule):
    def __init__(
        self,
        data_mean: float = 0,
        data_std: float = 1,
        learning_rate: float = 1e-3,
    ) -> None:
        super().__init__()
        self.data_mean = data_mean
        self.data_std = data_std
        self.learning_rate = learning_rate
        ...

    def training_step(self, batch):
        normalized_output = self(batch)
        denormalized_output = normalized_output * self.data_std + self.data_mean
        ...

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
        return optimizer

    ...

Here, learning_rate is an actual hyperparameter that one might want to set via a CLI argument. But data_mean and data_std are supposed to be calculated from the training data. These parameters are used to obtain a model that outputs values in real units, while the trainable parameters of the model are optimized to predict outputs with zero mean and unit standard deviation.

My main script looks something like this:

def main() -> None:
    cli = LightningCLI(MyModel, MyDataModule, run=False)

    data = MyDataModule("data")
    data.setup()
    model = MyModel(
        data_mean=data.train_mean.item(),
        data_std=data.train_std.item(),
    )
    trainer = L.Trainer()
    trainer.fit(model, data, ...)


if __name__ == "__main__":
    main()

Here, the idea would be to access the CLI arguments through the cli object and use them to implement my custom logic. But of course, the data_mean and data_std parameters end up in the config file and CLI help, where they don't belong.

Is there a way to exclude such parameters from the inferred CLI options or am I using LightningCLI wrong?

[mauvilsa]

The purpose of the config file that LightningCLI automatically saves is reproducibility. That is, it is for you to know what were the parameters that were used to train the model. If you define parameters such as data_mean: float, then they do belong in that config file.

Note that you have defined these parameters as non-optional. That means that every time an instance of MyModel is created, then these parameters always should get a float value. If not specified, they get the default, but still a float. If you want them to not always be a float, and implement some logic depending on that, then you could change the the type, e.g. to Optional[float]. All depends on what you want to do. But again, in the config saved during training, these parameters should appear just as they were given to train. If later, you want to do inference, and the values of these parameters should be different, then give the correct values. Regardless, that config is not intended to be directly for inference, since it purpose is only reproducibility.

[NiklasKappel]

The thing is I want the training to be reproducible, and the parameters should always get float values, but these values are not to be specified directly by the user (so they should not be exposed by the CLI interface). Instead, there is another non-optional entity the user must specify (here that's the train data set). Then there is a (pure) function that computes the actual model parameters from the options given by the user.

I guess I could make the model's __init__ take the train dataset as an argument and do the necessary computation by itself. Then it would take much more information than necessary, but at least it would be self-contained.

[mauvilsa]

A parameter that is in __init__ is something that should be specified by the user and not computed. If you compute something based on init params, then internally give it a different name since it is not the same as what gets provided.

What you need is argument linking. Have a look at #13403 (comment), also cli argument linking and jsonargparse argument linking. You could link the entire data module. Though, it might be cleaner to link a @property that provides whatever needs computing.

Note. When an argument is linked, it doesn't show up in the config. Since the value is derived instead of provided.

[NiklasKappel]

Got it, thanks! This is what I actually needed. :)