How can I use a weighted loss with LightningCLI?

[maxfreu]

I would like to use a weighted loss function in a LightningModule model. The weights can be computed after initializing the train dataset of the LightningDataModule, ergo I only have access to the weights after DataModule.setup(). However, LightningCLI only calls the respective init methods. I would need to setup the DataModule before initializing the model. So, what is the best solution to achieve this?

[mauvilsa]

You can use an argument link applied on instantiation, see argument-linking. This could be done in several ways, depending on when/where you want setup() to be called. One possibility could be that the LightningModule gets the instance of the LightningDataModule. This is something like:

class MyModule(LightningModule):
    def __init__(self, data_module: LightningDataModule):
        # data setup could be called here or later
        ...

class MyLightningCLI(LightningCLI):
    def add_arguments_to_parser(self, parser):
        parser.link_arguments("data", "model.data_module", apply_on="instantiate")

cli = MyLightningCLI(MyModule, MyDataModule)

Other possibilities are:

• The LightningDataModule implements a @property that automatically triggers setup and gives the data needed for the weighted loss. The LightningModule.__init__ has a parameter to receive this data. Then you create an instantiation link from the data property to the module parameter.
• The LightningModule.__init__ has a parameter to get the data needed for the weighted loss. Then create an instantiation link with source the LightningDataModule instance and a compute_fn function that takes care of setup() and returns the required data.

The link target would need .init_args. if the CLI uses subclass mode for the LightningModule, e.g., parser.link_arguments("data", "model.init_args.data_module", apply_on="instantiate").

[maxfreu]

Thank you! This worked. I opted for bullet point 1, as it seemed cleanest and easiest.