Lightning-AI/pytorch-lightning

Variable length number of batches is not supported

Open

#18.023 aperta il 8 lug 2023

Vedi su GitHub
 (2 commenti) (2 reazioni) (0 assegnatari)Python (3233 fork)batch import
bugdata handlinghelp wantedver: 2.0.x

Metriche repository

Star
 (26.687 star)
Metriche merge PR
 (Merge medio 9g 15h) (3 PR mergiate in 30 g)

Descrizione

Bug description

BatchSamplers that do not always return same number of batches are broken in Pytorch Lightning.

Related: https://github.com/Lightning-AI/lightning/issues/17793.

There are several issues:

  • batch_size is not required for a PyTorch Sampler--Pytorch Lightning is monkey patching this class such that a valid Sampler is incompatible with pytorch-lightning
  • even though Pytorch Lightning repeatedly calls len on the Sampler, this information is disregarded, and the epoch often terminates early
  • if you use DDP (>= 2 GPUs) and set the number of batches to a number larger than the actual number of batches (e.g. 100), it causes a deadlock
  • behavior of DDP and single-gpu differs (see output of batch_size)

The code below reproduces all of these issues.

What version are you seeing the problem on?

v2.0

How to reproduce the bug

import os

import torch, numpy as np
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer
from time import sleep


class RandomDataset(Dataset):
    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):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("train_loss", loss)
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("valid_loss", loss)

    def test_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("test_loss", loss)

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)

class VarLengthBatchSampler(torch.utils.data.sampler.Sampler):
    def __init__(self, data_source, max_batch_size=None, drop_last=True, seed=12340,
                 batch_size=None):
        """Randomly return batches of length [1,max_batch_size].
        
        Args:
            data_source: a dataset
            max_batch_size: the maximum batch size. This is required
            drop_last: if True, drop the last batch if too small
            seed: random seed
        """
        super().__init__(data_source)
        self.data_len = len(data_source)
        self.max_batch_size = max_batch_size
        self.seed = seed
        # in actual application, need to increment this in the training loop
        self.epoch = 0
        # self.approx_len = int(np.ceil(self.data_len / (self.max_batch_size // 2)))
        self.approx_len = 5 # to dramatically illustrate problem of not updating length
        # self.approx_len = 100 # causes deadlock!!
        print(f"approx_len: {self.approx_len}")
        self.len = self.approx_len # pytorch lightning
        self.drop_last = drop_last

    def rand_bz(self):
        return torch.randint(1,self.max_batch_size,(1,))[0]

    def __iter__(self):
        g = torch.Generator()
        g.manual_seed(self.seed + self.epoch)
        indices = torch.randperm(self.data_len, generator=g).tolist()
        batch = [] # accumulate idxs, then reset
        batches = [] # accumulate batches
        bz = self.rand_bz()
        # print(f"\n{bz=}")
        for idx in indices:
            batch.append(idx)
            if len(batch) == bz:
                batches.append(batch)
                bz = self.rand_bz()
                # print(f"\nnew {bz=}")
                batch = []
        if not self.drop_last and len(batch) > 0:
            batches.append(batch)
        
        # update length to match random sampling results
        self.len = len(batches)
        print(f"\nnumber of batches should be: {self.len}")
        return iter(batches)
    
    def __len__(self):
        print(f"\n__len__: {self.len=}")
        return self.len

class PLVarLengthBatchSampler(VarLengthBatchSampler):
    def __init__(self, data_source, batch_size=None, max_batch_size=8, drop_last=True, seed=12340):
        """Wrapper for for pytorch lightning compatibility
        """
        print(f"batch_size: {batch_size}")
        super().__init__(data_source, max_batch_size=max_batch_size, drop_last=drop_last, seed=seed)
        self.batch_size = max_batch_size

def run():
    train_dset = RandomDataset(32, 64)
    bz = 8
    train_data = DataLoader(train_dset,
        batch_sampler=PLVarLengthBatchSampler(train_dset, max_batch_size=bz))
    val_data = DataLoader(RandomDataset(32, 64), batch_size=32)
    test_data = DataLoader(RandomDataset(32, 64), batch_size=32)

    model = BoringModel()
    trainer = Trainer(
        # devices=1, # uses first __len__ call for all lengths
        devices=2, # uses approx_len for all lengths
        default_root_dir=os.getcwd(),
        num_sanity_val_steps=0,
        max_epochs=5,
        enable_model_summary=False
    )
    trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)
    trainer.test(model, dataloaders=test_data)

run()

Error messages and logs

One GPU

batch_size: None
approx_len: 5

number of batches should be: 15

__len__: self.len=15

__len__: self.len=15
Epoch 0:   0%|          | 0[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<?, ?it[/s](https://file+.vscode-resource.vscode-cdn.net/s)] 
__len__: self.len=15

number of batches should be: 16

__len__: self.len=16
Epoch 0:   7%|▋         | 1[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 327.83it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=16
Epoch 0:  13%|█▎        | 2[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 473.96it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=16
Epoch 0:  20%|██        | 3[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 558.74it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=16
Epoch 0:  27%|██▋       | 4[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 623.27it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=16
Epoch 0:  33%|███▎      | 5[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 634.23it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=16
Epoch 0:  40%|████      | 6[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 664.83it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=16
Epoch 0:  47%|████▋     | 7[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 695.54it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=16
Epoch 0:  53%|█████▎    | 8[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 721.10it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=16
Epoch 0:  60%|██████    | 9[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 739.82it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=16
Epoch 0:  67%|██████▋   | 10[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 756.82it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=16
Epoch 0:  73%|███████▎  | 11[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 771.80it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=16
Epoch 0:  80%|████████  | 12[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 785.72it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=16
Epoch 0:  87%|████████▋ | 13[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 791.61it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=16
Epoch 0:  93%|█████████▎| 14[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 799.92it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=16
Epoch 1:   0%|          | 0[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<?, ?it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]          
__len__: self.len=16

number of batches should be: 17

__len__: self.len=17
Epoch 1:   7%|▋         | 1[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 635.21it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=17
Epoch 1:  13%|█▎        | 2[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 756.34it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=17
Epoch 1:  20%|██        | 3[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 812.32it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=17
Epoch 1:  27%|██▋       | 4[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 838.86it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=17
Epoch 1:  33%|███▎      | 5[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 858.96it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=17
Epoch 1:  40%|████      | 6[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 862.05it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=17
Epoch 1:  47%|████▋     | 7[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 873.45it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=17
Epoch 1:  53%|█████▎    | 8[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 883.87it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=17
Epoch 1:  60%|██████    | 9[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 891.37it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=17
Epoch 1:  67%|██████▋   | 10[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 894.96it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=17
Epoch 1:  73%|███████▎  | 11[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 898.75it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=17
Epoch 1:  80%|████████  | 12[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 902.62it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=17
Epoch 1:  87%|████████▋ | 13[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 908.52it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=17
Epoch 1:  93%|█████████▎| 14[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<00:00, 914.53it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]
__len__: self.len=17
Epoch 2:   0%|          | 0[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00<?, ?it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=102]          
__len__: self.len=17

Two GPUs

LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]
LOCAL_RANK: 1 - CUDA_VISIBLE_DEVICES: [0,1]

number of batches should be: 22
number of batches should be: 22

batch_size: 8batch_size: 8

approx_len: 5
approx_len: 5
[/home/tyler/opt/anaconda/envs/gaddy/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:430](https://file+.vscode-resource.vscode-cdn.net/home/tyler/opt/anaconda/envs/gaddy/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:430): PossibleUserWarning: The dataloader, train_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 32 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
  rank_zero_warn(

__len__: self.len=5
__len__: self.len=5


__len__: self.len=5
__len__: self.len=5

[/home/tyler/opt/anaconda/envs/gaddy/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py:280](https://file+.vscode-resource.vscode-cdn.net/home/tyler/opt/anaconda/envs/gaddy/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py:280): PossibleUserWarning: The number of training batches (5) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.
  rank_zero_warn(
[/home/tyler/opt/anaconda/envs/gaddy/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:430](https://file+.vscode-resource.vscode-cdn.net/home/tyler/opt/anaconda/envs/gaddy/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:430): PossibleUserWarning: The dataloader, val_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 32 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
  rank_zero_warn(
Training: 0it [00:00, ?it[/s](https://file+.vscode-resource.vscode-cdn.net/s)]
Training:   0%|          | 0[/5](https://file+.vscode-resource.vscode-cdn.net/5) [00:00<?, ?it[/s](https://file+.vscode-resource.vscode-cdn.net/s)]
Epoch 0:   0%|          | 0[/5](https://file+.vscode-resource.vscode-cdn.net/5) [00:00<?, ?it[/s](https://file+.vscode-resource.vscode-cdn.net/s)] 
__len__: self.len=5
number of batches should be: 6


number of batches should be: 6

__len__: self.len=6
__len__: self.len=6

Epoch 0:  20%|██        | 1[/5](https://file+.vscode-resource.vscode-cdn.net/5) [00:00<00:01,  3.14it[/s](https://file+.vscode-resource.vscode-cdn.net/s)]
__len__: self.len=6
Epoch 0:  20%|██        | 1[/5](https://file+.vscode-resource.vscode-cdn.net/5) [00:00<00:01,  3.13it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=99]
__len__: self.len=6
Epoch 0:  40%|████      | 2[/5](https://file+.vscode-resource.vscode-cdn.net/5) [00:00<00:00,  6.22it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=99]
Epoch 0:  40%|████      | 2[/5](https://file+.vscode-resource.vscode-cdn.net/5) [00:00<00:00,  6.22it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=99]

__len__: self.len=6

__len__: self.len=6
Epoch 0:  60%|██████    | 3[/5](https://file+.vscode-resource.vscode-cdn.net/5) [00:00<00:00,  9.26it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=99]
__len__: self.len=6

__len__: self.len=6
Epoch 0:  80%|████████  | 4[/5](https://file+.vscode-resource.vscode-cdn.net/5) [00:00<00:00, 12.28it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=99]
__len__: self.len=6
Epoch 0: 100%|██████████| 5[/5](https://file+.vscode-resource.vscode-cdn.net/5) [00:00<00:00, 15.26it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=99]
[/home/tyler/opt/anaconda/envs/gaddy/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/logger_connector/result.py:432](https://file+.vscode-resource.vscode-cdn.net/home/tyler/opt/anaconda/envs/gaddy/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/logger_connector/result.py:432): PossibleUserWarning: It is recommended to use `self.log('valid_loss', ..., sync_dist=True)` when logging on epoch level in distributed setting to accumulate the metric across devices.
  warning_cache.warn(
Epoch 0: 100%|██████████| 5[/5](https://file+.vscode-resource.vscode-cdn.net/5) [00:00<00:00, 14.88it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=99]
__len__: self.len=6
Epoch 0:   0%|          | 0[/5](https://file+.vscode-resource.vscode-cdn.net/5) [00:00<?, ?it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=99]        
Epoch 1:   0%|          | 0[/5](https://file+.vscode-resource.vscode-cdn.net/5) [00:00<?, ?it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=99]

__len__: self.len=6
__len__: self.len=7


number of batches should be: 7

__len__: self.len=7

__len__: self.len=7
Epoch 1:  20%|██        | 1[/5](https://file+.vscode-resource.vscode-cdn.net/5) [00:00<00:00, 301.34it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=99]
__len__: self.len=7

__len__: self.len=7
Epoch 1:  40%|████      | 2[/5](https://file+.vscode-resource.vscode-cdn.net/5) [00:00<00:00, 380.35it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=99]
__len__: self.len=7
__len__: self.len=7

Epoch 1:  60%|██████    | 3[/5](https://file+.vscode-resource.vscode-cdn.net/5) [00:00<00:00, 424.97it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=99]
__len__: self.len=7
__len__: self.len=7

Epoch 1:  80%|████████  | 4[/5](https://file+.vscode-resource.vscode-cdn.net/5) [00:00<00:00, 441.45it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=99]
__len__: self.len=7
Epoch 1: 100%|██████████| 5[/5](https://file+.vscode-resource.vscode-cdn.net/5) [00:00<00:00, 290.58it[/s](https://file+.vscode-resource.vscode-cdn.net/s), v_num=99]
__len__: self.len=7

Environment

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):

More info

No response

cc @justusschock @awaelchli

Guida contributor

Direzione di ricerca
Analizza il codice di riproduzione per capire come Lightning gestisce i campionatori batch a lunghezza variabile, concentrandoti sul ciclo di caricamento dei dati e sulla sincronizzazione DDP.
Tech stack
python
Dominio
backend
Tipo issue
Bug
DifficoltàQuanto dovrebbe essere impegnativa per un nuovo contributor.
3
Tempo stimatoTempo stimato per completare e verificare uno scope piccolo.
Mezza giornata
Stato attivitàQuanto è probabile che la issue sia ancora attiva e reviewable.
Attiva
ChiarezzaQuanto chiaramente la issue descrive problema, scope e risultato atteso.
Chiara
Prerequisiti
PythonPyTorch
Adatta ai principiantiQuanto la issue è adatta a contributor alla prima esperienza.
65