Lightning-AI/pytorch-lightning

Variable length number of batches is not supported

Open

#18,023 建立於 2023年7月8日

在 GitHub 查看
 (2 留言) (2 反應) (0 負責人)Python (3,233 fork)batch import
bugdata handlinghelp wantedver: 2.0.x

倉庫指標

Star
 (26,687 star)
PR 合併指標
 (平均合併 9天 15小時) (30 天內合併 3 個 PR)

描述

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

貢獻者指南

研究方向
分析複現代碼，了解 Lightning 如何處理可變長度批次取樣器，重點關注資料載入迴圈和 DDP 同步。
技術棧
python
領域
backend
議題類型
錯誤
難度面向新貢獻者的預計實作難度，1 表示很小改動，5 表示專家級工作。
3
預計時間有經驗貢獻者完成調查、實作、測試並準備 pull request 的粗略時間範圍。
半天
活動狀態議題目前的可參與程度：新鮮、活躍、陳舊、阻塞或等待維護者輸入。
活躍
清晰度議題是否清楚說明預期改動、驗收標準和下一步。
清晰
前置要求
PythonPyTorch
新手友善度1-100 的估計分數，表示該議題對首次貢獻者的友善程度。
65