LightningCLI access DataModule methods in model constructor

[bstee615]

Hello, I am switching my code to use LightningCLI for easier config/reproducibility. I want to access dataset examples when creating my model, but I don't want to have to make my data loading code explicitly available to both LightningDataModule and LightningModule.

Currently my project uses code similar to this:

data = MyLightningDataModule(...)
model = MyLightningModule(data.input_dim, ...)
trainer.fit(model, datamodule=data)

where data.input_dim is roughly equivalent to data.train_dataloader().dataset[0].shape[0]. Basically, the width of the first layer of the model depends on the data generated by the DataModule. This is modeled after the API usage suggested in https://pytorch-lightning.readthedocs.io/en/stable/extensions/datamodules.html#using-a-datamodule. I posted pseudocode to make it simpler but I can post the real thing if I didn't get the point across.

This is how I use LightningCLI:

cli = LightningCLI(MyLightningModule, MyLightningDataModule)

But it seems that when using LightningCLI, the LightningModule is instantiated before the LightningDataModule, and I can't figure out how to configure LightningCLI to make the LightningDataModule available to the LightningModule constructor.

I've tried to make this happen by supplying datamodule as a parameter to my LightningModule following this API: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_cli.html#models-with-multiple-submodules.

class MyLightningModule(pl.LightningModule):
    def __init__(self, data: BigVulDatasetLineVDDataModule):
        self.data = data
        # instantiate model...

    def train_dataloader(self):
        return self.data.train_dataloader()

    def val_dataloader(self):
        return self.data.train_dataloader()

    def test_dataloader(self):
        return self.data.train_dataloader()

    # other LightningModule related methods...

And my config looks like this:

model:
  # model args...
  data:
    class_path: project.MyLightningDataModule
    init_args:
      # datamodule args...

...but it would be nice if I could keep these two classes independent. Is there any better way to achieve the functionality I want?

[mauvilsa]

For this you use link_arguments which requires to subclass LightningCLI. Something like:

class MyLightningModule(LightningModule):
    def __init__(self, input_dim: int):
        ...

class MyCLI(LightningCLI):
    def add_arguments_to_parser(self, parser):
        parser.link_arguments('data.input_dim', 'model.input_dim', apply_on='instantiate')

It might also be possible like your second example linking the entire data object. Though this needs the latest version of jsonargparse:

class MyLightningModule(pl.LightningModule):
    def __init__(self, data: BigVulDatasetLineVDDataModule):
        self.data = data
        # instantiate model...

class MyCLI(LightningCLI):
    def add_arguments_to_parser(self, parser):
        parser.link_arguments('data', 'model.data', apply_on='instantiate')

[bstee615]

This is what I needed! Thank you for pointing it out.

[OStars]

Thanks you for pointing link_arguments. But what if my model initialization needs parameters that were available after the datamodule was instantiated? Can the default LightningCli do this? It seems to only support initializing the model before initializing the datamodule

[terbed]

Yes, indeed! Here is the official documentation:
https://lightning.ai/docs/pytorch/stable/cli/lightning_cli_expert.html#argument-linking🎉

[calvinshopify]

Yes. Create a link applied on instantiation from data.label_num to the model. This should work as long as prepare_data gets run during data module init, or data.label_num is a @property that when accessed, runs prepare_data if not done already.

@mauvilsa this doesn't seem to be working as expected.

Here I do the linking to link max_event_id:

parser.link_arguments(
    "data.max_event_id",
    "model.embeddings.init_args.num_items",
    apply_on="instantiate",
)

Here is where I compute max_event_id (note that it is not included in __init__):

class SequenceDataModule(LightningDataModule):
    def __init__(
        self,
        ...
    ):
        self.max_event_id = get_max_event_id()
        log.info(f"Max event id: {self.max_event_id}")

Then I get the error "AttributeError: 'Namespace' object has no attribute 'max_event_id'".

[calvinshopify]

I have also tried:

        parser.link_arguments(
            "data",
            "model.embeddings.init_args.num_items",
            apply_on="instantiate",
            compute_fn=lambda data: data. get_max_event_id(),
        )

But that doesn't work because data is not instantiated (is a Namespace).

[mauvilsa]

Well, an error related to Namespace indicates that the class is not instantiated. But you have apply_on="instantiate" so this shouldn't be the case. There are unit tests that make sure this behavior is not broken as the code evolves. The only thing I can imagine is that something in your code makes it in some way different.

What can you do?

• Run export JSONARGPARSE_DEBUG=true before your script and see if in the logs give you a hint of what is the problem.
• Debug your code. For example, instead of compute_fn with a lambda make it a function and set a breakpoint inside of it.
• Copy here the stack trace.
• Tell us os, what versions of lightning and jsonargparse you are using.
• Provide to us a minimal reproduction script.

[calvinshopify]

Thanks for the response. The versions checked out (on latest of both Lightning and jsonarparse) ✅ and I created a smaller example to try and repro, but it worked as expected ✅

I was trying to port over a large Lightning implementation but started building it up incrementally since I got that weird behaviour and things seem to be working now. Not sure what the cause was but will circle back if I discover it.