Global GPU Flag

1

In older versions of PyTorch, in order to move everything to the GPU, one had to do the following.

# Define a lambda at the top
cuda = lambda x: x.cuda() if torch.cuda.is_available() else x

x = Variable(cuda(torch.randn(10))) # When creating variables
model = cuda(Model()) # When creating modules

With the release of PyTorch 0.4, this has been slightly simplified as:

# Define the default tensor type at the top
torch.set_default_tensor_type(torch.cuda.FloatTensor if torch.cuda.is_available() 
                                                     else torch.FloatTensor)

# When creating tensors
x = torch.randn(10) # No need for additional bells and whistles

However, while creating modules, there is no such simplification:

device = torch.device('cuda' if torch.cuda.is_avaliable() else 'cpu')
model = Model().to(device) # One has to still cast the module to the appropriate device

Is there any workaround possible?

2

Edit: This method doesn’t seem to work for reasons mentioned in a later post.

Just realized that you don’t actually need to cast modules to the device.

Setting the default tensor type applies to all parameters and tensors automatically within the module.

So, my current method for running a script on GPU is the following:

# Define this lambda
use_gpu = lambda x=True: torch.set_default_tensor_type(torch.cuda.FloatTensor 
                                             if torch.cuda.is_available() and x 
                                             else torch.FloatTensor)

Now, use this as a switch.

use_gpu()
''' ...
your code here
...'''

Thank you team for making life so simple :heart_eyes:

1 Like
3

Sorry to have revived this thread again.

Just tried it out and it seems that there are some fundamental problems with setting the default tensor type.

  • Dataloaders give normal (non-cuda) tensors by default. They have to be manually cast using the Tensor.to() method.

  • Many methods are simply not implemented for torch.cuda.*Tensor. Thus, setting the global tensor type to cuda fails.

  • Conversions to numpy using the numpy() method aren’t’ available for cuda tensors. One has to go x.cpu().numpy().
    Although this chain is agnostic, it defeats the purpose.

Can someone from the team pitch in on this?

Is there any hope for a global device setting that just works?