How does DataLoader pin_memory=True help with data loading speed?

Fei_Liu: <p>Suppose the original tensors in the dataset is not pinned. Then the <code>pin_memory=True</code> setting only (automatically) adds a pin operation for each of the tensors loaded from the dataset on the fly for that specific iteration (after collation, I believe). Isn’t that equivalent to these sequential steps?:</p>
<ol>
<li>pin an unpinned CPU tensor (which makes another copy to the non-pageable memory area).</li>
<li>(later in training) copy the pinned tensor to the corresponding GPU device.</li>
</ol>
<p>These two steps combined won’t be faster than directly copying an unpinned CPU tensor to GPU memory (it’s different from repetitively copying a pre-pinned CPU tensor to GPU, in which case it does save data transfer time). By this logic, the <code>pin_memory=True</code> option in DataLoader only adds some additional steps that are intrinsically sequential anyways, so how does it really help with data loading?</p>
<p>Note: this assumes <code>num_workers=0</code> (the default and most common case). If <code>num_workers</code> is not <code>0</code>, then the default process <code>fork</code> (I’m in linux) won’t work with different process pinning CPU tensors anyways (complaining about multiple CUDA initializations).</p>
<p>I think I must have some misunderstandings somewhere. Can someone help?</p>

ptrblck: <aside class="quote no-group" data-username="Fei_Liu" data-post="1" data-topic="181306">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/fei_liu/48/36219_2.png" class="avatar"> Fei_Liu:</div>
<blockquote>
<p>These two steps combined won’t be faster than directly copying an unpinned CPU tensor to GPU memory (it’s different from repetitively copying a pre-pinned CPU tensor to GPU, in which case it does save data transfer time). By this logic, the <code>pin_memory=True</code> option in DataLoader only adds some additional steps that are intrinsically sequential anyways, so how does it really help with data loading?</p>
</blockquote>
</aside>
<p>That’s not that case as pinned memory would allow you to trigger a non-blocking data transfer which won’t be possible otherwise. Take a look at <a href="https://github.com/pytorch/pytorch/issues/79933#issuecomment-1164055300">this description</a> I’ve posted in an issue for more details about the internals and how async copies are used.</p>
<aside class="quote no-group" data-username="Fei_Liu" data-post="1" data-topic="181306">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/fei_liu/48/36219_2.png" class="avatar"> Fei_Liu:</div>
<blockquote>
<p>Note: this assumes <code>num_workers=0</code> (the default and most common case). If <code>num_workers</code> is not <code>0</code>, then the default process <code>fork</code> (I’m in linux) won’t work with different process pinning CPU tensors anyways (complaining about multiple CUDA initializations).</p>
</blockquote>
</aside>
<p>Using pinned host memory won’t initialize a CUDA context and <code>num_workers&gt;0</code> in a <code>DataLoader</code> using <code>pin_memory=True</code> works.</p>

Fei_Liu: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="181306">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Using pinned host memory won’t initialize a CUDA context and <code>num_workers&gt;0</code> in a <code>DataLoader</code> using <code>pin_memory=True</code> works.</p>
</blockquote>
</aside>
<p>My bad. I re-tested it and you are right. I misconfigured something else that led to the previous error.</p>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="181306">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>That’s not that case as pinned memory would allow you to trigger a non-blocking data transfer which won’t be possible otherwise. Take a look at <a href="https://github.com/pytorch/pytorch/issues/79933#issuecomment-1164055300" rel="noopener nofollow ugc">this description</a> I’ve posted in an issue for more details about the internals and how async copies are used.</p>
</blockquote>
</aside>
<p>Follow up question on this: does it mean that if I don’t use non-blocking data transfer, then pin_memory won’t help? E.g. I only have one input tensor which is immediately used after <code>.to(device)</code> (hence non-blocking doesn’t help). Consider the following code snippet:</p>
<pre data-code-wrap="python"><code class="lang-python">data = torch.rand(1000, 1000, 1000)  # input data on CPU memory
dataset = TensorDataset()
dataloader = DataLoader(dataset, pin_memory=True)
for x, in dataloader:
  x = x.to('cuda:0', non_blocking=True)
  # do something with x immediately, e.g. feed it to some model on cuda:0
</code></pre>
<p>Is it fair to say that the <code>pin_memory=True</code> and <code>non_blocking=True</code> won’t help in this case? It only helps if:</p>
<ol>
<li>I do something else with CPU objects after the non_blocking transfer, or</li>
<li>I have multiple input data (e.g. <code>x, y, z, w ...</code> ) so that the host to device data transfer can be processed in a batch + non-blocking way (they are already pinned by the dataloader), or</li>
<li>I use multiple workers in dataloader (say <code>num_workers=4</code>) so that the pin_memory for multiple batch/iterations can happen simultaneously (not sure about this one, is it still sequential in the pin step?)</li>
</ol>
<p>Am I getting it?</p>

ptrblck: <aside class="quote no-group" data-username="Fei_Liu" data-post="3" data-topic="181306">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/fei_liu/48/36219_2.png" class="avatar"> Fei_Liu:</div>
<blockquote>
<p>Follow up question on this: does it mean that if I don’t use non-blocking data transfer, then pin_memory won’t help? E.g. I only have one input tensor which is immediately used after <code>.to(device)</code> (hence non-blocking doesn’t help).</p>
</blockquote>
</aside>
<p>Not quite as the CPU might be able to still run ahead and schedule the next kernel launch even if the kernel has to wait for the transfer to finish. Besides that you are right that an async transfer would allow you to overlap it with other work assuming it’s available.</p>