Torch，如何将多个 GPU 用于不同的数据集

Posted 2023-03-29

【中文标题】Torch，如何将多个 GPU 用于不同的数据集

【英文标题】：Torch, how to use Multiple GPU for different dataset

【发布时间】：2021-04-16 22:04:01

【问题描述】：

假设我有 4 个不同的数据集和 4 个 GPU，如下所示

4 个数据集

dat0 = [np.array(...)], dat1 = [np.array(...)] , dat2 = [np.array(...)] , dat3 = [np.array(...)] 

4 个显卡

device = [torch.device(f'cuda:i') for i in range(torch.cuda.device_count())]

假设四个数据集都已经转换成张量并转移到4个不同的GPU。

现在，我有一个来自其他模块的函数 f 可以在 GPU 上使用

如何同时进行以下操作，

计算这个的4个resulf

 ans0 = f(dat0) on device[0], ans1 = f(dat1) on device[1], ans2 = f(dat2) on device[2], ans3 = f(dat3) on device[3]

然后将所有4个ans移回cpu然后计算总和

 ans = ans0 + ans1 + ans2 + ans3   

【问题讨论】：

您真的打算一次性将四个数据集的全部发送到四个 GPU 吗？您打算在哪里计算最终总和？另外 - 我应该说非常重要 - 你打算从 ans 反向传播吗？

@Ivan 用于求和部分，它只是四个数组的求和并在 CPU 上完成。 ans 是我需要的最终结果

是的，但是你需要反向传播吗？

@Ivan 嗨，我不需要反向传播。谢谢你的回答，我试试看。

【参考方案1】：

假设您只需要ans 即可进行推理。您可以轻松地执行这些操作，但您肯定需要函数 f 同时在所有四个 GPU 上。

这是我会尝试的：复制 f 四次并发送到每个 GPU。然后计算中间结果，将每个结果发送回 CPU 进行最终操作：

fns = [f.clone().to(device) for device in devices]

results = []
for fn, data in zip(fns, datasets):
    result = fn(data).detach().cpu()
    results.append(result)

ans = torch.stack(results).sum(dim=0)

【讨论】：

一个问题是，它是顺序运行还是同时运行？

我认为不会，因为它不是异步代码。对于多 GPU 使用，您可能需要使用 nn.DataParallel