Lightning-AI/pytorch-lightning

Variable length number of batches is not supported

Open

#18.023 aberto em 8 de jul. de 2023

Ver no GitHub
 (2 comments) (2 reactions) (0 assignees)Python (3.233 forks)batch import
bugdata handlinghelp wantedver: 2.0.x

Métricas do repositório

Stars
 (26.687 stars)
Métricas de merge de PR
 (Mesclagem média 9d 15h) (3 fundiu PRs em 30d)

Description

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

Guia do colaborador

Direção de pesquisa
Analise o código de reprodução para entender como o Lightning lida com amostradores de lote de comprimento variável, concentrando se no loop de carregamento de dados e na sincronização DDP.
Pilha de tecnologia
python
Domain
backend
Tipo Issue
Bug
DifficultyDificuldade de implementação estimada para um novo contribuidor, de 1 para alterações muito pequenas a 5 para trabalho de nível especializado.
3
Tempo estimadoUm intervalo de tempo aproximado para um colaborador experiente investigar, implementar, testar e preparar um pull request.
Meio dia
Status da atividadeQuão disponível o issue aparece agora: novo, ativo, obsoleto, bloqueado ou aguardando entrada do mantenedor.
Ativo
ClarityCom que clareza o issue explica a mudança esperada, os critérios de aceitação e a próxima etapa.
Claro
Prerequisites
PythonPyTorch
Simpatia para novatosUma pontuação de 1 a 100 que estima o quão acessível este issue é para colaboradores iniciantes.
65