DataLoader multiprocessing with Dataset returning a CUDA tensor

Theo_Cheynel: <p>Hi,</p>
<p>I have a use case where my Dataset’s <code>__getitem__</code> fetches a Tensor somewhere in the memory of the main process, transfers it to the GPU to do some basic image processing faster than it ever could on the CPU, and returns a CUDA tensor directly. Then there’s a <code>DataLoader</code> on top of that.<br>
The <code>DataLoader</code> works fine when using <code>num_workers=0</code>, however I’m getting errors whenever I try to use multiprocessing by increasing <code>num_workers</code>…</p>
<p>Here’s a MWE :</p>
<pre><code class="lang-auto">import torch

class CudaDataset(torch.utils.data.Dataset):
    def __init__(self, device):
        self.tensor_on_ram = torch.Tensor([1, 2, 3])
        self.device = device

    def __len__(self):
        return len(self.tensor_on_ram)

    def __getitem__(self, index):
        return self.tensor_on_ram[index].to(self.device)

ds = CudaDataset(torch.device('cuda:0'))
dl = torch.utils.data.DataLoader(dataset=ds, batch_size=1, pin_memory=False, num_workers=2)

# First time runs with no issue at all
for i in dl:
    pass

## Let's do it a second time
for i in dl:  # Here it throws an error
    pass
</code></pre>
<p>Here’s the error :</p>
<pre><code class="lang-auto">RuntimeError: Caught RuntimeError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/usr/local/lib/python3.7/dist-packages/torch/utils/data/_utils/worker.py", line 287, in _worker_loop
    data = fetcher.fetch(index)
  File "/usr/local/lib/python3.7/dist-packages/torch/utils/data/_utils/fetch.py", line 49, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/usr/local/lib/python3.7/dist-packages/torch/utils/data/_utils/fetch.py", line 49, in &lt;listcomp&gt;
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "&lt;ipython-input-1-3f858a29a121&gt;", line 12, in __getitem__
    return self.tensor_on_ram[index].to(self.device)
  File "/usr/local/lib/python3.7/dist-packages/torch/cuda/__init__.py", line 207, in _lazy_init
    "Cannot re-initialize CUDA in forked subprocess. To use CUDA with "
RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method
</code></pre>
<p>I’ve tried adding the line <code>torch.multiprocessing.set_start_method('spawn')</code> at the top, but then I get <code>DataLoader worker (pid(s) 1078) exited unexpectedly</code></p>
<p>I’m not sure whether that use case is possible, I just wanted to benchmark the performance gain that I could obtain with more workers, as this is currently my bottleneck.</p>
<p>Does anyone know a way to work around this ? Thanks</p>

Ehsan1997: <p>I don’t think that moving the data to GPU right after loading it is a good idea.</p>
<p>An easy workaround would be to load the data in the getitem function and use the dataloader to fetch a batch of this data. Then you can move this data to GPU and apply the processing steps there.</p>
<p>Another workaround would be to pass a new collate_fn in the dataloader class.</p>

ejguan: <p>It’s better to use <code>pin_memory</code> in <code>DataLoader</code> to transfer data into shared memory. Then, move the Tensor to your main process to CUDA.</p>
<p>And, when you run <code>spawn</code>, please wrap your script using <code>if __name__ == "__main__":</code>.</p>

carbocation: <p>The advice in this thread seems to agree with the <a href="https://pytorch.org/docs/stable/data.html#multi-process-data-loading" rel="noopener nofollow ugc">pytorch documentation</a>, which I think is also worth linking to from here. E.g., they call out the following:</p>
<blockquote>
<p>It is generally not recommended to return CUDA tensors in multi-process loading because of many subtleties in using CUDA and sharing CUDA tensors in multiprocessing (see <a href="https://pytorch.org/docs/stable/notes/multiprocessing.html#multiprocessing-cuda-note" rel="noopener nofollow ugc">CUDA in multiprocessing</a>). Instead, we recommend using <a href="https://pytorch.org/docs/stable/data.html#memory-pinning" rel="noopener nofollow ugc">automatic memory pinning</a> (i.e., setting <code>pin_memory=True</code> ), which enables fast data transfer to CUDA-enabled GPUs.</p>
</blockquote>