LightningCLI: limit the output of known subclasses for better readability #20591
Description
Description & Motivation
The output of <subcommand> --help is populated with all known subclasses based on the docstring's type hints. As it can be seen in the following MWE, this can span multiple lines/pages in a terminal which IMHO limits readability.
It would be great if this output for subclasses could somehow be suppressed (either partially or completely).
# main.py
import lightning as L
import torch
from lightning.pytorch.cli import LightningCLI
class LitModel(L.LightningModule):
def __init__(self, model: torch.nn.Module):
super().__init__()
self.model = model
def forward(self, x):
return self.model(x)
def configure_optimizers(self):
return torch.optim.Adam(self.parameters())
def cli_main():
return LightningCLI(LitModel)
def cli_main():
cli = LightningCLI(LitModel)
# note: don't call fit!!
if __name__ == "__main__":
cli_main()
# note: it is good practice to implement the CLI in a function and call it in the main if blockpython main.py fit --help
And the output is:
...
--model.model MODEL (required, type: <class 'Module'>, known subclasses: torch.nn.Module, torch.nn.Identity, torch.nn.Linear, torch.nn.modules.linear.NonDynamicallyQuantizableLinear,
torch.nn.LazyLinear, torch.ao.nn.qat.Linear, torch.ao.nn.intrinsic.qat.LinearReLU, torch.ao.nn.qat.dynamic.Linear, torch.ao.nn.intrinsic.qat.LinearBn1d,
torch.ao.nn.quantized.reference.Linear, torch.nn.Bilinear, torch.nn.Threshold, torch.nn.ReLU, torch.ao.nn.quantized.ReLU6, torch.nn.RReLU, torch.nn.Hardtanh,
torch.nn.ReLU6, torch.nn.Sigmoid, torch.ao.nn.quantized.Sigmoid, torch.nn.Hardsigmoid, torch.nn.Tanh, torch.nn.SiLU, torch.nn.Mish, torch.nn.Hardswish,
torch.ao.nn.quantized.Hardswish, torch.nn.ELU, torch.ao.nn.quantized.ELU, torch.nn.CELU, torch.nn.SELU, torch.nn.GLU, torch.nn.GELU, torch.nn.Hardshrink,
torch.nn.LeakyReLU, torch.ao.nn.quantized.LeakyReLU, torch.nn.LogSigmoid, torch.nn.Softplus, torch.nn.Softshrink, torch.nn.MultiheadAttention,
(many more modules are shown here)
...
Pitch
No response
Alternatives
No response
Additional context
No response