使用cuda GPU的data.to（device）的使用

Question

我正在尝试实现一个神经网络来运行在8个GPU上，我只想澄清这些命令的确切作用

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
data.to(device)

这会自动将训练分散到8个GPU上吗？谢谢！

Answer 1

没有如果CUDA可用，代码段会将模型和数据移至GPU，否则，会将其放入CPU。

• [torch.device('cuda')指代当前的cuda设备
• [torch.device('cuda:0')指代带有index=0的cuda设备

要使用所有8个GPU，您可以执行以下操作：

if torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model)

注意：

• [torch.cuda.device_count()返回可用的GPU数量。
• 您无需致电：data = torch.nn.DataParallel(data)

为什么？因为torch.nn.DataParallel

是一个容器，它通过按批处理维度中的组块在指定的设备之间划分输入来并行化给定模块的应用程序（其他对象将每个设备复制一次）。在前向传递中，模块在每个设备上复制，每个副本处理输入的一部分。在向后传递过程中，来自每个副本的梯度被汇总到原始模块中。

批大小应大于使用的GPU数量。