Pin memory to specific GPU in DataLoader

LapoFrati: <p>Is there a way to specify the GPU to use when using a DataLoader with pin_memory = True?<br>
I’ve tried using<br>
<code>torch.cuda.set_device(device)</code><br>
and<br>
<code>with torch.cuda.device(gpu_id):</code><br>
<code>#DataLoader</code><br>
but nvidia-smi shows that the GPU used is always the first available.<br>
I think I could use:<br>
<code>CUDA_VISIBLE_DEVICES=x</code><br>
to make the target GPU be the only one available in the notebook but I’d rather not.</p>

SimonW: <p>afaik, pinned memory is not specific to any GPU. It just puts data in the locked page in host CPU memory. The destination is specified by .cuda(device)</p>

LapoFrati: <p>I thought that too but I’ve noticed that when I pin it, some stuff get loaded on the GPU according to nvidia-smi.<br>
I’ve tried changing the device before pinning it as in:<br>
<code>import torch</code><br>
<code>x_tensor = torch.FloatTensor(1)</code><br>
<code>with torch.cuda.device(2):</code><br>
<code>x_tensor.pin_memory()</code><br>
and I got this:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/8/8d3198e1a90b373fc3916ea1a4661945c5bf7def.png" data-download-href="https://discuss.pytorch.org/uploads/default/8d3198e1a90b373fc3916ea1a4661945c5bf7def" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/8/8d3198e1a90b373fc3916ea1a4661945c5bf7def_2_345x222.png" alt="image" data-base62-sha1="k93Ag5l2Whp7DT9Kq5z1TrAsCth" width="345" height="222" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/8/8d3198e1a90b373fc3916ea1a4661945c5bf7def_2_345x222.png, https://discuss.pytorch.org/uploads/default/optimized/2X/8/8d3198e1a90b373fc3916ea1a4661945c5bf7def_2_517x333.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/2X/8/8d3198e1a90b373fc3916ea1a4661945c5bf7def_2_690x444.png 2x" data-dominant-color="424142"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1266×818 73.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
I don’t know the inner working of pinning memory but this led me to believe some preparations take place on the GPU as well (setting the device before starting a DataLoader doesn’t seem to affect it).<br>
If I set the .cuda(device) on model/inputs they get sent to the correct GPU but that memory usage from pinning persists until I finish so I decided to ask on here for some opinions <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=12" title=":sweat_smile:" class="emoji" alt=":sweat_smile:" loading="lazy" width="20" height="20"></p>

SimonW: <p>I played with different sized tensors, the allocated memory is always constant. It is very likely that some context is created when you do so. So don’t worry too much about it.</p>

ngimel: <p>What pytorch version are you using? Pinning to specific GPU was merged recently <a href="https://github.com/pytorch/pytorch/pull/4196" rel="nofollow noopener">https://github.com/pytorch/pytorch/pull/4196</a>, before that context was always initialized on device 0.</p>

ngimel: <p>It is not such a small deal, on an 8 GPU system 7 unnecessary contexts can be created, each taking ~ 400 MB, together taking ~ 3 GB of memory.</p>

LapoFrati: <p>You are right <a class="mention" href="/u/ngimel">@ngimel</a> , my problem was that from conda I got the latest release 0.3.0 but the docs I was reading were already about version 0.4.0 <img src="//discuss.pytorch.org/images/emoji/apple/expressionless.png?v=5" title=":expressionless:" class="emoji" alt=":expressionless:"></p>
<p>This commit should fix my problem:<br>
<aside class="onebox githubcommit">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/commit/bc6bd62bd6805136117d860819a361ef78608462#diff-724910f4a1ffa1852a02e149b181ac22" target="_blank" rel="nofollow noopener">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
      <a href="https://github.com/csarofeen" target="_blank" rel="nofollow noopener">
    <img alt="csarofeen" src="https://avatars2.githubusercontent.com/u/22205833?v=4" class="thumbnail onebox-avatar" width="90" height="90">
  </a>

<h4>
  <a href="https://github.com/pytorch/pytorch/commit/bc6bd62bd6805136117d860819a361ef78608462" target="_blank" rel="nofollow noopener">Fix distributed dataloader so it pins memory to current GPU not GPU 0.</a>
</h4>


<div class="date">
  by <a href="https://github.com/csarofeen" target="_blank" rel="nofollow noopener">csarofeen</a>
  on <a href="https://github.com/pytorch/pytorch/commit/bc6bd62bd6805136117d860819a361ef78608462" target="_blank" rel="nofollow noopener">06:58PM - 15 Dec 17</a>
</div>

<div class="github-commit-stats">
  changed <strong>1 files</strong>
  with <strong>7 additions</strong>
  and <strong>3 deletions</strong>.
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>

pickandroll3: <p>I am having the same issue with <code>torch-1.10.2+cu113</code>, is this fixed?</p>

ptrblck: <p>Yes, the initialization of the CUDA context on the default device when pinned memory is used, was fixed in late 2017.</p>