Best practice for re-using arguments between multiple parent classes in pl cli

[Jimmy2027]

Hello, I'm wondering what the recommended practice is when using LightningCLI and reusing multiple arguments between multiple parent classes.
In the example below, all arguments to instantiate the baseclass BaseModel have to be defined in both ImageClassifier1 and ImageClassifier2, together with their default values which can get cumbersome for many parameters and confusing when mixing up defaults.

Is there a better way to do this?

Example:

class BaseModel(LightningModule):
    def __init__(self, model, base_arg: bool = False, lr, gamma, batch_size):
        super().__init__()
        self.base_arg = base_arg
        self.model = model
        self.test_acc = Accuracy()

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self.forward(x)
        return F.nll_loss(logits, y.long())

    def configure_optimizers(self):
        optimizer = torch.optim.Adadelta(self.model.parameters(), lr=self.hparams.lr)
        return [optimizer], [torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=self.hparams.gamma)]


@MODEL_REGISTRY
class ImageClassifier1(BaseModel):
    def __init__(self, arg1: int = 1, lr=1.0, gamma=0.7, batch_size=32):
        model = Net1(arg1=arg1)
        super().__init__(model=model, lr=lr, gamma=gamma, batch_size=batch_size)
        self.save_hyperparameters()


@MODEL_REGISTRY
class ImageClassifier2(BaseModel):
    def __init__(self, arg2: int = 1, lr=1.0, gamma=0.7, batch_size=32):
        model = Net2(arg2=arg2)
        super().__init__(model=model, lr=lr, gamma=gamma, batch_size=batch_size)
        self.save_hyperparameters()

[rohitgr7]

cc: @carmocca

[tchaton]

Hey @Jimmy2027,

You could register them directly as follows:

MODEL_REGISTRY(ImageClassifier1, key="key to remember the arguments of ImageClassifier1")
MODEL_REGISTRY(ImageClassifier2, key="key to remember the arguments ImageClassifier2")

[Jimmy2027]

hi @tchaton, thanks for you reply :) I'm not sure how to use the key argument in this context, but I found creating a dataclass storing default values to be a good option:

@dataclass
class ModuleBaseArgs:
    lr: float = 1.0
    gamma: float = 0.7

@MODEL_REGISTRY
class ImageClassifier1(BaseModel):
    def __init__(self, arg1: int = 1, lr=ModuleBaseArgs.lr, gamma=ModuleBaseArgs.gamma):
        model = Net1(arg1=arg1)
        super().__init__(model=model, base_arg=True)
        self.save_hyperparameters(ignore="model")


@MODEL_REGISTRY
class ImageClassifier2(BaseModel):
    def __init__(self, arg2: int = 1, lr=ModuleBaseArgs.lr, gamma=ModuleBaseArgs.gamma):
        model = Net2(arg2=arg2)
        super().__init__(model=model, base_arg=True)
        self.save_hyperparameters(ignore="model")

[carmocca]

This structure is perfectly fine. Remember that you can compose multiple configuration files

python script.py fit --config=model_optimization_args.yaml --config=model_data_args.yaml

You could also use a configuration object that's shared between classes or use Mixins but they are not "better" solutions, just different ways of setting up the structure.

[Jimmy2027]

I see, thanks for your answer :)

Is there any way to store in the config which model one is running?
I guess this could be done by setting a default which_model argument like this:

@MODEL_REGISTRY
class ImageClassifier1(BaseModel):
    def __init__(self, arg1: int = 1, lr=ModuleBaseArgs.lr, gamma=ModuleBaseArgs.gamma, which_model:str = 'imgclf1'):
        model = Net1(arg1=arg1)
        super().__init__(model=model, base_arg=True)
        self.save_hyperparameters(ignore="model")


@MODEL_REGISTRY
class ImageClassifier2(BaseModel):
    def __init__(self, arg2: int = 1, lr=ModuleBaseArgs.lr, gamma=ModuleBaseArgs.gamma, which_model:str = 'imgclf2'):
        model = Net2(arg2=arg2)
        super().__init__(model=model, base_arg=True)
        self.save_hyperparameters(ignore="model")

or is there a better way to do it?

[carmocca]

If you do python script.py fit --model=ImageClassifier2 --model.arg2=123 --print_config

You should see:

model:
    class_path: __main__.ImageClassifier2
    init_args:
        arg2: 123
        ...

So the class path acts like the label you are looking for

[Jimmy2027]

ah very nice! Is there also a way to obtain this information in the config, without any command line arguments?

[carmocca]

When the config file is generated, it reflects whichever configuration is set. If the model class is not part of your arguments then it won't show.