Torch，如何为不同的数据集使用多个 GPU

Question

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

4 个数据集

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

4 GPU

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

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

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

如何同时进行以下操作，

计算这个的 4 个结果

 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个答案移回cpu然后计算总和

 ans = ans0 + ans1 + ans2 + ans3

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

Torch，如何为不同的数据集使用多个 GPU

Torch, how to use Multiple GPU for different dataset

torch

pytorch

torchvision