[utils.bottleneck] Bottleneck crashes with multi-threaded data loader

[fmassa]

torch.utils.bottleneck doesn't work properly when the code contains a data loader that uses more than 0 threads.

Minimum reproducible example (mwe.py):

import argparse
import torch
import torch.utils.data

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='mwe')
    parser.add_argument('--num-workers', default=0, type=int)
    args = parser.parse_args()

    data = torch.rand(10, 1000)
    target = torch.rand(10)
    dataset = torch.utils.data.TensorDataset(data, target)
    data_loader = torch.utils.data.DataLoader(dataset,
        batch_size=2, num_workers=args.num_workers)
    for i, batch in enumerate(data_loader):
        pass

Running the script via:

python -m torch.utils.bottleneck -- mwe.py --num-workers 0

works fine, while

python -m torch.utils.bottleneck -- mwe2.py --num-workers 1

crashes with the following stack trace:

Traceback (most recent call last):
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/bottleneck/__main__.py", line 280, in <module>
    main()
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/bottleneck/__main__.py", line 261, in main
    autograd_prof_cpu, autograd_prof_cuda = run_autograd_prof(code, globs)
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/bottleneck/__main__.py", line 155, in run_autograd_prof
    result.append(run_prof(use_cuda=True))
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/bottleneck/__main__.py", line 149, in run_prof
    exec(code, globs, None)
  File "mwe2.py", line 15, in <module>
    for i, batch in enumerate(data_loader):
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 285, in __next__
    return self._process_next_batch(batch)
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 306, in _process_next_batch
    raise batch.exc_type(batch.exc_msg)
RuntimeError: Traceback (most recent call last):
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 57, in _worker_loop
    samples = collate_fn([dataset[i] for i in batch_indices])
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 57, in <listcomp>
    samples = collate_fn([dataset[i] for i in batch_indices])
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/data/dataset.py", line 40, in __getitem__
    return tuple(tensor[index] for tensor in self.tensors)
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/data/dataset.py", line 40, in <genexpr>
    return tuple(tensor[index] for tensor in self.tensors)
RuntimeError: /private/home/fmassa/github/pytorch/torch/csrc/autograd/profiler.h:52: initialization error

assigning this to @zou3519 , even thought I'm not sure if it's a problem in the profiler or in the bottleneck tool.

pytorch version '0.4.0a0+b21e135'

cc @ezyang @gchanan @zou3519 @ssnl

[zou3519]

Yikes, thanks for the catch. I'll look into it

[TiRune]

Ran into this problem as well on unix; Would be great if we could still use the bottleneck functionality with multiple dataloaders :)

[jhagege]

Hi, has anyone found any workaround with this issue ?
I am trying to identify bottleneck in a PyTorch code and have a hard time doing it without proper profiling.
Can anyone recommend good resources for identifying bottleneck (even with other techniques) ?
Thanks !

[stmharry]

This thread deserves more visibility.

Tried many ways to avoid /pytorch/torch/csrc/autograd/profiler.h:52: initialization error, and ended up here, just to realize the problem is the number of workers.

Have not been able to locate a good tool in identifying the bottleneck.

[ezyang]

@stmharry A very simple bandaid which I'd be happy to approve, is to update the "initialization error" message with a suggested workaround of reducing the number of worker threads. Would you like to submit a PR?

[dcela]

This is still an issue for me as well

[yaceben]

Still an issue for me as well