How does one make sure that everything is running on GPU automatically in Pytorch?

Question

I wanted a way with minimal amount of code such that everything in my script runs automatically in GPU (or the standard way pytorch did it). Something like:

torch.everything_to_gpu()

and then it "just works". I don't care about manually putting things in GPU etc. I just want it to do its stuff automatically (sort of the way tensorflow does it?). I did see a related question in the pytorch forum but it doesn't seem that they address my issue directly.

Right now it seems to me (from the examples I've been through) that one can do something like what I want by specifying a simple type to every torch Variable/tensor as follows:

dtype = torch.FloatTensor
# dtype = torch.cuda.FloatTensor # Uncomment this to run on GPU

so as long as every variable/tensor is takes dtype somehow eg

Variable(torch.FloatTensor(x).type(dtype), requires_grad=False)

then we can use that single variable to control what is in GPU and not. The issue that I am encountering that makes things ambiguous to me if such a single command exists is when using torch.nn.Module package. For example when using

l = torch.nn.Linear(D_in,D_out)

or costume NN classes (that inherit from it). Such cases seems that the best way to deal with it is to use the:

torch.nn.Module.cuda(device_id=device_id) # device_id = None is the default

function/method. However this seems to suggest to me that there might be other hidden functions that I might not be aware of to make sure that everything does indeed run in GPU.

Thus: Is there a centralized way to make sure everything runs in some (ideally automatically) assigned GPU?

---

In reflection I think one thing that is confusing me is that I don't understand the model of how pytorch carriers on computations on GPU. For example, I am fairly certain tht the way MATLAB works is that if at least one thing is on GPU then all further computations will be on GPU. So I guess, I am wondering, is this how pytorch works? If possible, how does it compare to TensorFlow?

Answer 1

I think that there is no such thing.

From what I've seen, people usually create classes that:
i) Inherit from nn.Module .
ii) Have an attribute describing model parameters (eg self.opt );
iii) Set each variable/parameters as attributes (eg self.my_var )
iv) Then call .cuda() on it if a kind of -use_gpu parameter is set.

I also use a maybe_cuda(variable) function inside my classes in order to create variables easier (pass a Variable, return variable.cuda() if opt.cuda is True.

In fact, I did something like this (may not be perfect, but found it practical):

class MyModule(nn.Module):
    def __init__(self, opt):
        super(MyModule, self).__init__()
        self.opt = opt

    def maybe_cuda(self, variable):
        if self.opt.cuda:
            return variable.cuda()
        return variable

class Model(MyModule):
    def __init__(self, opt, other_arg):
        super(Model, self).__init__(opt)

        self.linear = nn.Linear(opt.size1, opt.size2)
        self.W_out = nn.Parameter(_____)

    def forward(self, ____):
        # create a variable, put it on GPU if possible
        my_var = self.maybe_cuda(Variable(torch.zeros(___)))

Answer 2

As far as I know, there isn't a config that makes all future torch related computations happen in the GPU. You need to specify it explicitly.

In addition to pltrdy's answer , the following is another easy way to run your code in the GPU. Create a torch.device variable that has the device information on the device you want your code to run on.

#Assuming you have a cuda boolean flag
device = torch.device('cuda') if cuda else torch.device('cpu')
some_tensor = torch.randn(2).to(device)
some_model = ModelClass().to(device)

#If you want to use multiple GPUs use DataParallel
some_model = torch.nn.DataParallel(some_model)

As far your last question, new tensors dependent on a tensor will automatically reside in the parent tensor's device.

a = torch.randn(1).cuda() #In the GPU
b = a + 2                 #Also in the GPU

Documentation for DataParallel

Answer 3

Seems to be an open issue on github: [FR] torch context with default device & dtype . I came here because I figured it was a dumb question and there is good reason not to do it (here [feature request] Global GPU Flag they mention talk about undebuggable scenarios where you have a program running tensors on multiple devices...). It seems like it's something enough people want and they may implement it.