Access datamodule in setup()

[gianmarcoaversanoenx]

Hello,
I want to dynamically build a model. One model attribute will be inferred when the data is available. Thus, I need to access the dataset that is being used for the training from within the LightningModule. I tried to access that from the setup() method, but I failed.

def setup(self, stage: str = None) -> None:
        """Called at the beginning of fit (train + validate), validate, test, or predict. This is a good hook when you need to build models dynamically or adjust something about them. This hook is called on every process when using DDP.
        Args:
            stage (Optional[str]): either 'fit', 'validate', 'test', or 'predict'
        """
        if stage == "fit" and self.model.num_nodes_pmf is None:
            loaders: DataLoader = self.trainer.train_dataloader.loaders  # type: ignore
            graphs_train = loaders.dataset
            tmp_gg = [len(to_networkx(gg).nodes) for gg in graphs_train]  # type: ignore
            self.model.num_nodes_pmf = np.bincount(tmp_gg)[1:]

Error message:

if stage == "fit" and self.model.num_nodes_pmf is None:
>           loaders: DataLoader = self.trainer.train_dataloader.loaders  # type: ignore
E           AttributeError: 'NoneType' object has no attribute 'loaders'

What would be the correct way to do this?

[tshu-w]

May you can access datamodule throughself.trainer.datamodule ?

[gianmarcoaversanoenx]

I must not. Generally, you do not know how the user has passed the data (datamodule or simply train_dataloader via the Trainer).

In the end, I solved as follows:

    def on_train_start(self) -> None:
        """Called at the beginning of training after sanity check."""
        if self.model.num_nodes_pmf is None:
            loaders: DataLoader = self.trainer.train_dataloader.loaders  # type: ignore
            graphs_train = loaders.dataset
            tmp_gg = [len(to_networkx(gg).nodes) for gg in graphs_train]  # type: ignore
            self.model.num_nodes_pmf = torch.from_numpy(np.bincount(tmp_gg)[1:])

But I do not understand why the loaders are not available in the setup() method. This method is supposed to allow users to dynamically set up their model. Dynamically should mean "from the data", when this will be available. Instead, I had to use a hook.