What if I load data in __init__ function of LightningDataModule

[zhiqiangdon]

Hi,

I see the doc describing the functions of LightningDataModule.

https://pytorch-lightning.readthedocs.io/en/latest/notebooks/lightning_examples/datamodules.html#Defining-The-MNISTDataModule

Here is my thinking. If some variable, e.g., a transform, can be defined in init function, and later shared across different GPUs. Theoretically, if we load data in init, the data should also be able to transfer to different GPUs similarly. In the case of a single machine with multiple GPUs, the data will be copied multiple times in the memory. In the case of multiple machines, the data will broadcast through the network from the main node to other nodes. Broadcasting large data through networks may have efficiency issue, which is why we had better load data in the setup function.

Please let me know whether my analysis is correct or not. Basically, I am not clear about how the variables, e.g., i.e. self.something, defined in init are shared across multiple GPUs/machines. Thanks!

[tchaton]

@zhiqiangdon,

The data don't get broadcasted. They are just manipulated (downloaded, loaded, processed, etc...)

Loading isn't any issue, but downloading or preprocessing in a distributed setting is.

Imagine the following scenario with 2 processes.

A -> download the data at time t1 and B start to download the data at time t2. Both process will override each other and the data will corrupted. Same thing with preprocessing.

If you don't do any data manipulation, it is entirely fine loading the data in init.

However, if you do, use the prepare_data method called only on a single process (rank 0) and anything else you want within the setup method.

[zhiqiangdon]

Thanks @tchaton,

1. Based on your description, I guess prepare_data has some lock mechanism in the distributed setting. That is, if one process starts running prepare_data, it will block others from running prepare_data to avoid duplication. Right?
2. If using multiple machines, I guess using prepare_data may become tricky. There are two cases. First, if multiple machines share the same disk, only one process among all the machines should run prepare_data. Otherwise, the data overriding and corruption may happen. Second, if each machine has an independent disk, then each machine should have one process running prepare_data. Is this the underhood logic of how prepare_data works?
3. Do you mean each process runs __init__ function independently in a similar way as setup?
4. If __init__ and setup functions both run distributedly, is one reason splitting them due to the prepare_data? The running order is __init__ -> prepare_data -> setup. If there is no need to download and preprocess data, then we can move code from setup to __init__. Right?

[rohitgr7]

@zhiqiangdon

1. lightning just runs prepare_data on the main process before the distributed process actually starts so there is no blocking happening behind the scenes.
2. To tackle this issue we have prepare_data_per_node. A node is just a machine. If they share the disk then prepare_data_per_node should be set to False.
3. User runs the __init__ function when they initialize the DataModule, lightning just send to across devices.

[zhiqiangdon]

Thanks @rohitgr7,

Just to clarify. If I use self.data = load_data(data_path) in __init__, will the data be loaded only by the rank 0 process? If so, will the self.data be sent to across devices later by lightning?

[tchaton]

Hey @zhiqiangdon,

__init__ data will be called on processes along with self.data = load_data(data_path) if you implemented in such a way.

[zhiqiangdon]

Got it. Thanks @tchaton!

[tarokiritani]

@tchaton I am still confused. Suppose I have code like below:

datamodule = MyLightningDataModule() # __init__() is run.
model = MyLightningModule()
trainer = Trainer(gpus=[0, 1]) # use ddp.
trainer.fit(model, datamodule) # prepare_data() is run in ONE process.

If data are manipulated somehow in __init__() or prepare_data(), the manipulation happens only once, assuming that trainer does not do anything magical. I still do not understand what should be included in __init__() vs. prepare_data().