Dynamically Setting out_dim Based on LightningDataModule in LightningCLI

[aditya0by0]

Dynamically Setting out_dim Based on LightningDataModule in LightningCLI

Hi all,

I’m looking for a way to dynamically set the out_dim parameter in my LightningModule. Currently, this is passed explicitly via the command line, but I want to avoid that and instead have it automatically set based on the num_of_labels from the LightningDataModule. The problem is that num_of_labels is determined inside the setup method of the DataModule, not during its initialization.

I tried using link_arguments to link the num_of_labels from the DataModule to the out_dim in the LightningModule, but it doesn’t seem to work for this use case.

I’ve created a minimal example of the code below to illustrate the issue:
Note : num_of_labels can only be determined in setup method of the datamodule in below case, and not in constructor.

import torch
from lightning.pytorch import LightningDataModule, LightningModule
import torch.nn.functional as F
from lightning.pytorch.cli import LightningCLI
from torch.utils.data import DataLoader, TensorDataset


class MyLightningDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self._num_of_labels = None
        self._feature_vector_size = None

    def prepare_data(self):
        pass

    def setup(self, stage=None):
        self._num_of_labels = 10
        self._feature_vector_size = 20

    @property
    def num_of_labels(self):
        return self._num_of_labels

    @property
    def feature_vector_size(self):
        return self._feature_vector_size

    def train_dataloader(self):
        assert self.feature_vector_size is not None, "feature_vector_size must be set"
        # Dummy dataset for example purposes
        x = torch.randn(100, self.feature_vector_size)
        y = torch.randn(100, self._num_of_labels)
        dataset = TensorDataset(x, y)
        return DataLoader(dataset, batch_size=32)


class MyLightningModel(LightningModule):
    def __init__(self, input_dim: int, out_dim: int):
        super().__init__()
        self.save_hyperparameters()
        assert out_dim is not None, "out_dim must be set"
        assert input_dim is not None, "input_dim must be set"
        # self.fc1 = torch.nn.Linear(20, 10)
        self.fc1 = torch.nn.Linear(input_dim, out_dim)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

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

    def training_step(self, batch, batch_idx):
        x, y = batch
        loss = F.mse_loss(self.forward(x), y)
        self.log("train_loss", loss)
        return loss


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

        parser.link_arguments(
            "data.num_of_labels", "trainer.callbacks.init_args.num_labels", apply_on="instantiate"
        )

        # Further I need `num_of_labels` for `torchmetrics.MetricCollection`, commented as of now
        # command: fit --model.train_metrics="path_to_yaml_file"
        # for kind in ("train", "val", "test"):
        #     for average in ("micro-f1", "macro-f1"):
        #         parser.link_arguments(
        #             'data.num_of_labels',
        #             f"model.init_args.{kind}_metrics.init_args.metrics.{average}.init_args.num_labels",
        #             apply_on="instantiate"
        #         )


if __name__ == '__main__':
    cli = MyCLI(MyLightningModel, MyLightningDataModule,)

As you can see, I’m trying to use link_arguments to link num_of_labels to out_dim, but since num_of_labels is determined in the setup method (not the constructor), I’m running into issues.

Has anyone encountered a similar situation or found a way to dynamically set the out_dim based on the data module? Is there a better way to achieve this, or am I missing something?

Any feedback or suggestions would be greatly appreciated!

Thanks in advance!

Environment

• PyTorch Lightning: 2.1.2
• Python: 3.10.14
• Torch: 2.5.1
• jsonargparse: 4.31.0

[mauvilsa]

You have two options:

1. [preferred] In the data module, change the setup method so that if it is called multiple time, only the first run does the required logic. Then change the num_of_labels property so that it calls setup before accessing the _num_of_labels attribute.
2. Create a link applied on instantiation with source the entire data instance (instead of an attribute of it) and a compute_fn that first runs setup and then accesses the num_of_labels property.