question: Training slows down significantly for small dataset sizes
context: 🐛 Bug
  When the dataset size is small (i.e. comparable to the minibatch size), it slows training down significantly.
  No GPU, batch size 64, dataset size 1024: 185 iterations/second
  No GPU, batch size 64, dataset size 100: 47 iterations/second
  1 GPU, batch size 64, dataset size 1024: 110 iterations/second
  1 GPU, batch size 64, dataset size 100: 23 iterations/second
  1 GPU, batch size 800, dataset size 1024: 19 iterations/second
  1 GPU, batch size 800, dataset size 10000: 90 iterations/second
  1 GPU, batch size 64, dataset size 10000: 235 iterations/second
  Please reproduce using the BoringModel
  import os, sys
  from argparse import ArgumentParser
  
  import torch
  from torch.utils.data import Dataset, DistributedSampler, DataLoader
  
  from pl_examples import cli_lightning_logo
  from pytorch_lightning import LightningModule, Trainer
  from pytorch_lightning.utilities.seed import seed_everything
  from pytorch_lightning.callbacks.progress import ProgressBar, ProgressBarBase, tqdm, reset, convert_inf
  
  class CustomProgressBar(ProgressBar):
      def init_train_tqdm(self) -> tqdm:
          """ Override this to customize the tqdm bar for training. """
          bar = tqdm(
              desc='Training',
              initial=self.trainer.global_step,
              position=(2 * self.process_position),
              disable=self.is_disabled,
              leave=True,
              dynamic_ncols=True,
              file=sys.stdout,
              smoothing=0,
          )
          return bar
      def on_train_start(self, trainer, pl_module):
          super(ProgressBar, self).on_train_start(trainer, pl_module)
          self.main_progress_bar = self.init_train_tqdm()
          self.prev_train_gs = -1
          reset(self.main_progress_bar, self.trainer.max_steps)
  
      def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
          super(ProgressBar, self).on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx)
          if self.prev_train_gs != self.trainer.global_step and self._should_update(self.trainer.global_step, self.trainer.max_steps):
              self._update_bar(self.main_progress_bar)
              self.main_progress_bar.set_postfix(trainer.progress_bar_dict)
              self.prev_train_gs = self.trainer.global_step
  
      def on_train_epoch_start(self, trainer, pl_module):
          super(ProgressBar, self).on_train_epoch_start(trainer, pl_module)
  
      def on_train_end(self, trainer, pl_module):
          super(ProgressBar, self).on_train_end(trainer, pl_module)
  
  class RandomDataset(Dataset):
      """
      >>> RandomDataset(size=10, length=20)  # doctest: +ELLIPSIS
      <...bug_report_model.RandomDataset object at ...>
      """
  
      def __init__(self, size, length):
          self.len = length
          self.data = torch.randn(length, size)
  
      def __getitem__(self, index):
          return self.data[index]
  
      def __len__(self):
          return self.len
  
  
  class BoringModel(LightningModule):
      """
      >>> BoringModel()  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
      BoringModel(
        (layer): Linear(...)
      )
      """
  
      def __init__(self, train_data, test_data, bs):
          """
          Testing PL Module
  
          Use as follows:
          - subclass
          - modify the behavior for what you want
  
          class TestModel(BaseTestModel):
              def training_step(...):
                  # do your own thing
  
          or:
  
          model = BaseTestModel()
          model.training_epoch_end = None
  
          """
          super().__init__()
          self.layer1 = torch.nn.Linear(32, 32)
          self.layer2 = torch.nn.Linear(32, 32)
          self.layer3 = torch.nn.Linear(32, 2)
  
          self.train_data = train_data
          self.test_data = test_data
          self.bs = bs
  
      def forward(self, x):
          return self.layer3(torch.relu(self.layer2(torch.relu(self.layer1(x)))))
  
      def loss(self, batch, prediction):
          # An arbitrary loss to have a loss that updates the model weights during `Trainer.fit` calls
          return torch.nn.functional.mse_loss(prediction, torch.ones_like(prediction))
  
      def step(self, x):
          x = self.forward(x)
          out = torch.nn.functional.mse_loss(x, torch.ones_like(x))
          return out
  
      def training_step(self, batch, batch_idx):
          output = self.forward(batch)
          loss = self.loss(batch, output)
          return {"loss": loss}
  
      def training_step_end(self, training_step_outputs):
          return training_step_outputs
  
      def training_epoch_end(self, outputs) -> None:
          torch.stack([x["loss"] for x in outputs]).mean()
  
      def configure_optimizers(self):
          optimizer = torch.optim.Adam(list(self.layer1.parameters()) + list(self.layer2.parameters()) + list(self.layer3.parameters()), lr=0.001)
          lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
          return [optimizer], [lr_scheduler]
  
      def train_dataloader(self):
          train_loader = DataLoader(self.train_data, shuffle=True, num_workers=1, batch_size=self.bs)
          return train_loader
  
  parser = ArgumentParser()
  parser.add_argument("--gpus", type=int, default=0)
  parser.add_argument("--num_processes", type=int, default=1)
  parser.add_argument("--dataset_size", type=int, default=1024)
  parser.add_argument("--mb_size", type=int, default=64)
  args = parser.parse_args()
  
  
  def test_run():
      # data
      train_data = torch.randn(args.dataset_size, 32)
      test_data = torch.randn(256, 32)
  
      # model
      model = BoringModel(train_data, test_data, bs=args.mb_size)
      trainer = Trainer(
          gpus=args.gpus,
          logger=False,
          max_steps=5000,
          limit_val_batches=0,
          num_processes=args.num_processes,
          weights_summary=None,
          reload_dataloaders_every_epoch=False,
          callbacks=[CustomProgressBar()]
      )
  
      # fit
      trainer.fit(model)
  
      print(f"{trainer.accelerator_backend=}")
      print(f"{trainer.gpus=}")
      print(f"{trainer.num_processes=}")
      print(f"{trainer.global_step=}")
  
  if __name__ == "__main__":
      test_run()
  To Reproduce
  Run the following command: python bug_report.py --gpus 1 --dataset_size 10000 --mb_size 64
  for varying values of gpus, dataset_size, and mb_size.
  Expected behavior
  Iterations/second is unaffected by dataset size.
  Environment
  
  PyTorch Version (e.g., 1.0): 1.8.1
  OS (e.g., Linux): Linux
  How you installed PyTorch (conda, pip, source): conda
  Python version: 3.9
  
  Additional context
  My guess is that this is caused by inter-epoch reloading of the dataset. The code should be restructured to pre-load a fixed number of minibatches ahead, rather than caring about the location of epoch boundaries.
answer: I would say that the case here is that you would need to assume always some overhead and for a small dataset the initial phase is dominant compare to the full run, you can see a parallel with a car riding 100 or 1000 meters, in both cases, you need to start from zero and as long you go you benefit from no need starting again...