machine-learning - 如何确保 Pytorch 中的所有内容都自动在 GPU 上运行？

Question

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

torch.everything_to_gpu()

然后它就“正常工作”了。我不关心手动将东西放入 GPU 等。我只是希望它自动完成它的事情(有点像 tensorflow 那样？)。我确实看到了a related question in the pytorch forum但他们似乎并没有直接解决我的问题。

现在在我看来(从我经历过的例子来看)，人们可以通过为每个 torch 变量/张量指定一个简单类型来完成我想要的事情，如下所示:

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

只要每个变量/张量都以某种方式采用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 中运行。

因此:是否有一种集中的方式来确保一切在某些(最好是自动)分配的 GPU 中运行？

<小时/>

经过反射(reflection)，我认为让我困惑的一件事是我不理解 pytorch 如何在 GPU 上进行计算的模型。例如，我相当确定 MATLAB 的工作方式是，如果至少一件事是在 GPU 上进行的，那么所有进一步的计算都将在 GPU 上进行。所以我想，我想知道，这就是 pytorch 的工作原理吗？如果可能的话，它与 TensorFlow 相比如何？

Answer 1

据我所知，没有一个配置可以使所有 future 与 torch 相关的计算都在 GPU 中进行。您需要明确指定它。

除了 pltrdy's answer ，以下是在 GPU 中运行代码的另一种简单方法。创建一个 torch.device 变量，其中包含您希望运行代码的设备的设备信息。

#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)

就您的最后一个问题而言，依赖于张量的新张量将自动驻留在父张量的设备中。

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

DataParallel 的文档