gaozhihan/PreDiff

About Using Other Datasets for Training

qwertyui20 opened this issue · 5 comments

When training with the Sevirlr dataset, the model needs to pass in the CATALOG.csv file and use the already defined data loader SEVIRDataLoader. If training with another dataset, how should CATALOG.csv and the data loader be handled? Do we need to redefine it?

Thanks for your interest. The most convenient way to use your own data is to implement a LightningDataModule and integrate it into your training script. Here is an example to illustrate how to do it.

  1. Implement your own LightningDataModule (example link). You may refer to its official doc for more detailed instructions.
  2. Integrate your new LightningDataModule into your training script (example link).
  3. Modify the data layout in the forward() method of the training LightningModule (example link).

Thank you very much for your answer! I have implemented a DataModule for my own data, which can correctly load data on the Earthformer model. I plan to use this DataModule in the Prediff as well. Regarding the third step you provided, could you provide more suggestions? Do I need to modify the forward and more functions?
Currently, my Dataloader seems to be invalid. Could you provide some suggestions?
屏幕截图 2024-04-10 110836

I noticed you mentioned Earthformer. Do you reuse any of the code directly from the Earthformer project? If so, there may be a compatibility issue between the versions of pytorch_lightning and lightning used. pytorch_lightning==1.6 used in Earthformer is not compatible with lightning==2.0 used in PreDiff. It's important to ensure everything is using lightning==2.0 here, including LightningDataModule, to avoid any potential bugs from version mismatches.

Hello, I have completed the training of VAE and alignment stages. I noticed that in the vae and alignment stages, the dimensions of the data in the data loader are all three-dimensional. However, under the premise of using the same data loader, there is an error in the number of dimensions in prediff. Is the data used in the prediff stage the same as Are the data used in the other two stages different?

The data layouts in the dataloaders for VAE (N1HWC) and PreDiff (NTHWC) training are the same. The data layouts when fed into VAE (NCHW) and PreDiff (NTHWC) are different.
You could find the exact information about the data layout for VAE in:

target_bchw = rearrange(batch, "b 1 h w c -> b c h w").contiguous()

and for PreDiff in:

target_seq_bchw = rearrange(target_seq, "b t h w c -> (b t) c h w")