How to loop over the entire dataset before training start

[dalessioluca]

Hello,
my LightiningModule has an imputer which needs to be initialized from data (think about mean-imputation) before the training start.

Is there any ModelHook which allows me to perform a single pass on the entire dataset before the training start ?

Thanks in advance

class Net(pl.LightningModule):

def __init__(self):
      self.imputer = MyImputer()   # the imputer need to be initialized from the data
      self.net = MyNet()

def training_step(self, bath, batch_idx):
      x, y_hat = batch
      x_imputed = self.imputer(x)
      y = self.net(x_imputed)
      loss = (y-y_hat).pow(2).sum()
      return loss


def before_training(self, batch, batch_idx):
     # How can I do this? Is there any ModelHook I can use/modify?
     # I need to perform a full pass of the data to initialize the imputer before the training begins
     self.imputer.initialize_from_data(batch)

[rohitgr7]

you can do

def on_train_start(self):
    for batch in self.trainer.train_dataloader.loaders:
        ...

[dalessioluca]

Brilliant! Thanks!