Lightning-AI / Lightning-AI/pytorch-lightning

Unable to run custom sampler with custom arguments on TPU

Open
#17,262 0 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

accelerator: tpu bug data handling
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

I'm trying to implement the meta-learning model following this [tutorial](https://lightning.ai/docs/pytorch/2.0.1/notebooks/course_UvA-DL/12-meta-learning.html). The tutorial uses a custom sampler with new arguments (e.g. N_way, K_shot, etc).

When I run the training on Colab TPU, it throws an error in this [file](https://github.com/Lightning-AI/lightning/blob/master/src/lightning/pytorch/utilities/data.py#L295). Looks like the new arguments cause the error because lightning expects the custom sampler has the same arguments as BatchSampler.

How can we deal with this?

### How to reproduce the bug

```python
https://colab.research.google.com/drive/1zZpHWE33ZkivrdfuoClZCp6cm2cJlKIP#scrollTo=c89b030b&line=3&uniqifier=1
```

### Error messages and logs

```
INFO: [rank: 0] Received SIGTERM: 15
File "/usr/local/lib/python3.9/dist-packages/lightning/pytorch/strategies/launchers/xla.py", line 107, in _wrapping_function
results = function(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/lightning/pytorch/trainer/connectors/data_connector.py", line 487, in _process_dataloader
dataloader = trainer._data_connector._prepare_dataloader(dataloader, shuffle=is_shuffled, mode=stage)
File "/usr/local/lib/python3.9/dist-packages/lightning/pytorch/loops/evaluation_loop.py", line 169, in setup_data
dl = _process_dataloader(trainer, dl)
File "/usr/local/lib/python3.9/dist-packages/lightning/pytorch/trainer/connectors/data_connector.py", line 200, in _prepare_dataloader
dataloader = _update_dataloader(dataloader, sampler, mode=mode)
File "/usr/local/lib/python3.9/dist-packages/lightning/pytorch/utilities/data.py", line 132, in _update_dataloader
dl_args, dl_kwargs = _get_dataloader_init_args_and_kwargs(dataloader, sampler, mode)
TypeError: __init__() missing 1 required positional argument: 'K_shot'
```

### Environment

Current 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 @carmocca @JackCaoG @steventk-g @Liyang90 @justusschock @awaelchli

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with lightning/pytorch/utilities/data.py around line 295 and reproduce the failure using the linked Colab TPU example and meta-learning tutorial. Trace how the dataloader is reconstructed when the custom sampler uses N_way and K_shot; done means the custom sampler runs on TPU without the missing K_shot initialization error.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.