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

user_1601580 0
Viewed 2596

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 e.g.

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?

conv-neural-network pytorch machine-learning neural-network
3 Answers

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

user_65230630

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.

user_60845170
Related