Repository: pytorch
Issue: Cannot re-initialize CUDA in forked subprocess
Labels: oncall: distributed,module: dataloader,module: cuda,triaged

## 🐛 Bug

I am getting CUDA re-initialize error.  
I am using below code to generate synthetic dataset on GPU. To perform distributed training I am using [official PyTorch distributed training helper](https://github.com/pytorch/pytorch/blob/master/torch/distributed/launch.py) script.
```
import random
import torch
import torch.utils.data as data


class DummyDataset(data.Dataset):
    def __init__(self, num_classes):
        super(DummyDataset, self).__init__()
        # Create tensor on GPU 
        self.tensor = torch.randn(3, 224, 224, device=torch.device('cuda'))
        self.num_classes = num_classes

    def __getitem__(self, index):
        torch.manual_seed(index)
        random.seed(index)
        return self.tensor, \
               random.randint(0, self.num_classes - 1)
```
It throws following error 

```
in execute
    class_targets) in enumerate(train_dataset_loader):
  File "/usr/local/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 345, in __next__
    data = self._next_data()
  File "/usr/local/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 856, in _next_data
    return self._process_data(data)
  File "/usr/local/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 881, in _process_data
    data.reraise()
  File "/usr/local/lib/python3.6/site-packages/torch/_utils.py", line 394, in reraise
    raise self.exc_type(msg)
RuntimeError: Caught RuntimeError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/usr/local/lib/python3.6/site-packages/torch/utils/data/_utils/worker.py", line 178, in _worker_loop
    data = fetcher.fetch(index)
  File "/usr/local/lib/python3.6/site-packages/torch/utils/data/_utils/fetch.py", line 47, in fetch
    return self.collate_fn(data)
  File "/usr/local/lib/python3.6/site-packages/torch/utils/data/_utils/collate.py", line 79, in default_collate
    return [default_collate(samples) for samples in transposed]
  File "/usr/local/lib/python3.6/site-packages/torch/utils/data/_utils/collate.py", line 79, in <listcomp>
    return [default_collate(samples) for samples in transposed]
  File "/usr/local/lib/python3.6/site-packages/torch/utils/data/_utils/collate.py", line 53, in default_collate
    storage = elem.storage()._new_shared(numel)
  File "/usr/local/lib/python3.6/site-packages/torch/storage.py", line 126, in _new_shared
    return cls(size)
  File "/usr/local/lib/python3.6/site-packages/torch/cuda/__init__.py", line 477, in _lazy_new
    _lazy_init()
  File "/usr/local/lib/python3.6/site-packages/torch/cuda/__init__.py", line 195, in _lazy_init
    "Cannot re-initialize CUDA in forked subprocess. " + msg)
RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

```
I tried to follow the workaround mentioned in this [discussion](https://discuss.pytorch.org/t/runtimeerror-cannot-re-initialize-cuda-in-forked-subprocess-to-use-cuda-with-multiprocessing-you-must-use-the-spawn-start-method/14083/6) and changed my `DummyDataSet` class as follows

```
import random
import torch
torch.multiprocessing.set_start_method('spawn')# good solution !!!!
import torch.utils.data as data


class DummyDataset(data.Dataset):
    def __init__(self, num_classes):
        super(DummyDataset, self).__init__()
        # Create tensor on GPU 
        self.tensor = torch.randn(3, 224, 224, device=torch.device('cuda'))
        self.num_classes = num_classes

    def __getitem__(self, index):
        time.sleep(0.1);# !!!!!! In order to test, should be have virtual process time !!!!!!
        torch.manual_seed(index)
        random.seed(index)
        return self.tensor, \
               random.randint(0, self.num_classes - 1)
```
This gives another error 
```
  File "/pytorch_sample/dummy_data.py", line 3, in <module>
    torch.multiprocessing.set_start_method('spawn')# good solution !!!!
  File "/usr/local/lib/python3.6/multiprocessing/context.py", line 242, in set_start_method
    raise RuntimeError('context has already been set')
RuntimeError: context has already been set
Traceback (most recent call last):
  File "/usr/local/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 761, in _try_get_data
    data = self._data_queue.get(timeout=timeout)
  File "/usr/local/lib/python3.6/queue.py", line 173, in get
    self.not_empty.wait(remaining)
  File "/usr/local/lib/python3.6/threading.py", line 299, in wait
    gotit = waiter.acquire(True, timeout)
  File "/usr/local/lib/python3.6/site-packages/torch/utils/data/_utils/signal_handling.py", line 66, in handler
    _error_if_any_worker_fails()
RuntimeError: DataLoader worker (pid 1108) exited unexpectedly with exit code 1. Details are lost due to multiprocessing. Rerunning with num_workers=0 may give better error trace.
````

I am 
## Environment
```
Collecting environment information...
PyTorch version: 1.4.0a0
Is debug build: No
CUDA used to build PyTorch: 10.0

OS: Ubuntu 16.04.6 LTS
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.12) 5.4.0 20160609
CMake version: version 3.5.1

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 10.0.130
GPU models and configuration: 
GPU 0: Tesla V100-SXM2-32GB
GPU 1: Tesla V100-SXM2-32GB
GPU 2: Tesla V100-SXM2-32GB
GPU 3: Tesla V100-SXM2-32GB
GPU 4: Tesla V100-SXM2-32GB
GPU 5: Tesla V100-SXM2-32GB
GPU 6: Tesla V100-SXM2-32GB
GPU 7: Tesla V100-SXM2-32GB

Nvidia driver version: 440.33.01
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5

Versions of relevant libraries:
[pip3] numpy==1.18.4
[pip3] torch==1.4.0a0
[pip3] torchvision==0.5.0a0
[conda] Could not collect
```


## Additional context

<!-- Add any other context about the problem here. -->


cc @ngimel @SsnL @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @xush6528 @osalpekar