How to make sure everything runs on the GPU automatically in Pytorch?
I want a way with as little code as possible so that everything in my script runs automatically on the GPU (or the standard Pytorch way did it). Something like:
torch.everything_to_gpu()
and then it "just works". I don't care about manually putting things on the GPU, etc. I just want it to do its stuff automatically (how does the tensor process do it?). I saw a related question on the pytorch forum , but they don't seem to address my issue directly.
Right now it seems to me (from the examples I went through) that it is possible to do something like what I want by specifying a simple type on each Variable / tensor torch like this:
dtype = torch.FloatTensor
# dtype = torch.cuda.FloatTensor # Uncomment this to run on GPU
as long as each variable / tensor takes dtype
somehow for example.
Variable(torch.FloatTensor(x).type(dtype), requires_grad=False)
then we can use that single variable to control what is in the GPU instead of. The problem I am facing makes things ambiguous for me if such a single command exists when using the package torch.nn.Module
. For example, when using
l = torch.nn.Linear(D_in,D_out)
or suits of NN classes (which inherit from it). Cases like this show that the best way to deal with it is to use:
torch.nn.Module.cuda(device_id=device_id) # device_id = None is the default
function / method. However, this seems to be telling me that there may be other hidden features that I may not be aware of, to make sure everything is actually GPU powered.
Thus: is there a centralized way to make sure everything is running on some (ideally automatic) GPU assignment?
In reflection, I think one thing that is confusing me is that I don't understand the model of how pytorch media do GPU computing. For example, I'm fairly confident in how MATLAB works: if at least one is on the GPU, then all further calculations will be on the GPU. So I guess I'm wondering, is this the way pytorch works? If possible, how does it compare to TensorFlow?
source to share
I think 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 the parameters of the model (for example self.opt
);
iii) set each variable / parameters as attributes (for example self.my_var
)
iv) then call .cuda()
on it if any parameter is given -use_gpu
.
I also use a function maybe_cuda(variable)
inside my classes to make it easier to create a variable (pass the variable, return variable.cuda()
if opt.cuda
True.
I actually did something like this (maybe not ideal, 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(___)))
source to share