How to get reference of data module in model, or access model in data module

[ElderWanng]

For simplicity, I want to separate the model and data module. But my dataset is a resample dataset(fairseq provide an example). It's an auxiliary dataset for multitask learning. It downsamples a subset of a huge dataset. I write a set_epoch method, for resetting the subset of dataset joining training before every beginning of the epoch. If not so, the resample dataset returns same samples every epochs. Currently, I did it by

  def on_train_epoch_start(self) -> None:
      self.train_dataset.set_epoch(self.trainer.current_epoch)

Now I want to separate the model part and data part. Is there any way to access the data module in the model? Because the data module doesn't provide on_train_epoch_start hook.

Same way, I'm want to let two-part access each other. Because tokenizer could be used in two parts.

[rohitgr7]

LightningModule does provide on_train_epoch_start hook.

anyways you can access it using trainer.datamodule

also:

self.train_dataset.set_epoch(self.trainer.current_epoch)

this is done inside fit_loop by default.