How to make torch.cuda.device context managers compatible with cpu

username: oneilllml
cooked: <p>I’ve been working on a machine with multiple GPUs, and have needed to specific which GPU to run on. <a href="https://pytorch.org/docs/stable/notes/cuda.html" rel="nofollow noopener">The docs</a> say that best practice is to use a <code>torch.cuda.device</code> context manager, so the following line appears at multiple points in my code:</p>
<pre><code class="lang-auto">with torch.cuda.device(f"cuda:{ARGS.gpu}"):
    &lt;do-tensor-stuff&gt;
</code></pre>
<p>Now I want to make this code able to run on a CPU. Understandably, this line is now giving the error</p>
<pre><code class="lang-auto">AssertionError: 
Found no NVIDIA driver on your system. Please check that you
have an NVIDIA GPU and installed a driver from
http://www.nvidia.com/Download/index.aspx
</code></pre>
<p>I can’t think of a neat solution. Obviously the following is not good because there would be so much duplicate code.</p>
<pre><code class="lang-auto">if torch.cuda.is_available():
    with torch.cuda.device(f"cuda:{ARGS.gpu}"):
        &lt;do-tensor-stuff&gt;
else:
    &lt;do-tensor-stuff&gt;
</code></pre>

username: ptrblck
cooked: <p>Alternatively to the context manager you could also create a <code>device</code> variable as explained in the <a href="https://pytorch.org/docs/stable/notes/cuda.html#device-agnostic-code">device-agnostic code part</a>:</p>
<pre><code class="lang-python">args = parser.parse_args()
args.device = None
if not args.disable_cuda and torch.cuda.is_available():
    args.device = torch.device(f"cuda:{ARGS.gpu}")
else:
    args.device = torch.device('cpu')
</code></pre>
<p>Afterwards just pass the <code>args.device</code> to each <code>to()</code> operation.</p>

username: oneilllml
cooked: <p>Thanks. Yes this works, but what would be ideal is if the context manager could take a cpu device as an argument. Then we could avoid having to call <code>.to(device)</code> on every new tensor.</p>

username: ptrblck
cooked: <p>The current context manager isn’t working like that for CUDA tensors either, is it?<br>
You would still need to call <code>to()</code>, <code>device=</code> or <code>cuda()</code>. It will just change the default device:</p>
<pre><code class="lang-python">with torch.cuda.device(1):
    # allocates a tensor on GPU 1
    a = torch.tensor([1., 2.], device=cuda)

    # transfers a tensor from CPU to GPU 1
    b = torch.tensor([1., 2.]).cuda()
    # a.device and b.device are device(type='cuda', index=1)

    # You can also use ``Tensor.to`` to transfer a tensor:
    b2 = torch.tensor([1., 2.]).to(device=cuda)
    # b.device and b2.device are device(type='cuda', index=1)
</code></pre>

username: oneilllml
cooked: <p>Yes that’s true. Actually, the case I have in mind involves tensor-creating functions.</p>
<p>I used to pass a device argument to all such functions, which became cumbersome when there were a lot and they were nested and things.</p>
<pre><code class="lang-auto">def make_tensor_device(device): 
    return torch.tensor([1.,2.]).to(device)
   
device = torch.device(f"cuda:{ARGS.device}")
t = make_tensor_device(device)
</code></pre>
<p>The context manager simplified this because I can just set the default gpu and then call <code>.cuda()</code> on created tensors.</p>
<pre><code class="lang-auto">def make_tensor():
    return torch.tensor([1.,2.]).cuda() 

device = torch.device(f"cuda:{ARGS.device}")  
with torch.cuda.device(device):
    t = make_tensor()
</code></pre>
<p>Now it seems I have to go back to using the device argument. Not the end of the world, but it would be nice to allow setting the default device to either the cpu or a specific gpu, maybe that would even mean that <code>.cuda()</code> would no longer be necessary.</p>

username: ptrblck
cooked: <p>Well, there is the option of setting the default tensor type via:</p>
<pre><code class="lang-python">torch.set_default_tensor_type(torch.cuda.FloatTensor)
</code></pre>
<p>but I would recommend avoiding it, as this will create all tensors on the device, which are not necessarily needed on the GPU and you might run in a lot of issues.</p>
<p>I understand that the <code>device</code> argument might be a bit cumbersome, but I personally think it’s the cleanest way of writing the code.</p>
<p>However, any suggestions on improving the user experience are more then welcome.<br>
As already said, I’m not a huge fan of changing the default dtype, but how would a “good” API look like for your use case?<br>
Could you post a pseudo code for it?</p>

username: oneilllml
cooked: <p>I was hoping to be able to do something like:</p>
<pre><code class="lang-auto">def make_tensor():
    return torch.tensor([1.,2.])

if not ARGS.disable_cuda and torch.cuda.is_available():
    device = torch.device(f"cuda:{ARGS.gpu}")
else:
    device = torch.device('cpu')

with torch.device(device):
    t = make_tensor() 
</code></pre>
<p>and have this return a tensor on either the cpu or the specified gpu. However, this might also have the problem of putting lots of unnecessary things on the gpu. I didn’t know about the option to set the default type, and if you think that’s not a good idea, then maybe what I’m suggesting isn’t either.</p>

username: oneilllml
cooked: <p>Hmm maybe it would work if <code>.cuda()</code> just did nothing when the user had set the device to ‘cpu’. The following kind of works to give what I’m after.</p>
<pre><code class="lang-auto">from contextlib import contextmanager                                           
import torch                                                                          

torch.manually_set_device='none'                                                   
                                                                                     
@contextmanager                                                                    
def set_torch_device(new_device):                                                  
    prev_device = torch.manually_set_device                                        
    torch.manually_set_device = new_device                                         
    yield                                                                          
    torch.manually_set_device = prev_device                                        
                                                                                     
def maybe_cuda(t):                                                                 
    if torch.manually_set_device == 'none':                                        
        return t.cuda()                                                            
    else:                                                                          
        return t.to(torch.manually_set_device)                                     
                                                                                     
def make_tensor():                                                                 
    return maybe_cuda(torch.Tensor([1.,2.]))                                       
  
device = torch.device('cpu')                                                    
with set_torch_device(device):                                                                                                                                                                    
    t = make_tensor()                                                              
    print(t.device)                                                                
                                                                                     
t = make_tensor()                                                                  
print(t.device)       
</code></pre>

username: googlebot
cooked: <p>Note that for all tensor creation functions without explicit values (empty,zero,full,randn etc.) host-device transfer (.cuda()) is much slower than device=X argument to initial function. Also, I think you can change manually_set_device from ‘none’ to None and avoid if-else checks. I.e.</p>
<pre><code class="lang-auto">cur_device = None
#@contextmanager doing DI
...
def make_tensor(shape):
  return torch.zeros(shape, device=cur_device)
</code></pre>