Home Articles

如何确保Pytorch中的所有内容都在GPU上自动运行?

Asked
Viewed 158 times
3

我想用一种代码量最少的方法,这样我脚本中的所有东西都可以在GPU中自动运行(或者pytorch的标准方式就是这样) . 就像是:

torch.everything_to_gpu()

然后它"just works" . 我不关心手动将东西放在GPU等中 . 我只是想让它自动完成它的东西(有点像tensorflow那样吗?) . 我确实看到了a related question in the pytorch forum但似乎并没有直接解决我的问题 .

现在,在我看来(从我经历过的例子中)可以通过为每个火炬变量/张量指定一个简单类型来实现我想要的东西,如下所示:

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

所以只要 every 变量/张量以某种方式采取 dtype ,例如

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

然后我们可以使用该单个变量来控制GPU中的内容而不是 . 如果存在这样的单个命令,那么我遇到的问题会使我的模糊不清是在使用 torch.nn.Module 包时 . 例如使用时

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

或服装NN类(继承自它) . 这种情况似乎是处理它的最好方法是使用:

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

功能/方法 . 然而,这似乎告诉我,可能还有其他隐藏的功能,我可能不知道确保一切确实在GPU中运行 .

因此:是否有一种集中的方法来确保 everything 在某些(理想情况下自动)分配的GPU中运行?


在反思中,我认为让我感到困惑的一件事是,我不理解pytorch如何在GPU上进行计算的模型 . 例如,我相当确定MATLAB的工作方式是,如果至少有一件事在GPU上,那么所有进一步的计算将在GPU上 . 所以我想,我想知道,这是pytorch的工作原理吗?如果可能,它与TensorFlow相比如何?

1 Answer

  • 2

    我认为没有这样的事情 .

    从我所见,人们通常会创建以下类:
    i)继承自 nn.Module .
    ii)具有描述模型参数的属性(例如 self.opt );
    iii)将每个变量/参数设置为属性(例如 self.my_var
    iv)如果设置了一种 -use_gpu 参数,则在其上调用 .cuda() .

    我也在我的类中使用 maybe_cuda(variable) 函数,以便更容易地创建变量(传递变量,如果 opt.cuda 为True,则返回 variable.cuda() .

    事实上,我做了这样的事情(可能不完美,但发现它很实用):

    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(___)))
    

Related