Is pytorch lightning slow between epochs

[pamparana34]

I am using pytorch lightning and during the epoch I am getting around 25 iterations/second, which is comparable to vanilla pytorch code.

However, the time taken between an epoch ending and the next one starting seems to be relatively large, order of few seconds. Is this normal behaviour on the side of pytorch lightning. This really massively increases my training time.

So, I am creating my train_dataloaderas follows:

def train_dataloader(self):
        weighted_sampler = WeightedRandomSampler(
            weights=self.dataset.weights,
            num_samples=len(self.dataset.weights),
            replacement=True,
        )
        return DataLoader(
            self.dataset,
            shuffle=False,
            sampler=weighted_sampler,
            batch_size=self.batch_size,
            num_workers=4,
            pin_memory=True,
            persistent_workers=True
        )

The persistent_workers flag seems important else my number of iterations severely drops.

My train step is nothing special:

def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
        neg_ELBO, _, _, _ = self.forward(batch)
        self.log("train/loss", neg_ELBO, on_step=False, on_epoch=True, prog_bar=False)
        return {"loss": neg_ELBO}

My dataset does some heavy lifting in its constructor but this is only created once as far as I can tell. I am at a loss as to why the delay would be so pronounced. To put this into perspective, my whole epoch takes a second but to start the next epoch takes another 3....

Is there anyway I can get a hint of what might be happening? Or is there some tricks I can use to speed this up?

[rahulvigneswaran]

Hey @pamparana34

1. Do you have any custom def training_epoch_end?
2. Are you using ddp/dp? If ddp, how many node and gpus?
3. Did you try increasing the num_workers?

[pamparana34]

1: No custom training_epoch_end.
2: No dip/dp. Just single A5000 GPU.
3: Increasing the num_workers does not have any affect on this behaviour. Experimented with a range of 1 - 32.

[akihironitta]

@pamparana34 You can use the profiler to get some insight into your code.

trainer = Trainer(profiler="simple")
trainer = Trainer(profiler="advanced")

https://pytorch-lightning.readthedocs.io/en/1.5.10/common/trainer.html#profiler