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

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

4数据集

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

4 GPU

假设所有四个数据集都已转换为张量并传输到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个AN移回cpu，然后计算总和

 ans = ans0 + ans1 + ans2 + ans3   

假设您只需要

ans

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

f

下面是我要尝试的：复制

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吗？您计划在哪里计算最终金额？另外，我应该说，非常重要的是，您是否计划从

ans

反向传播？@Ivan对于求和部分，它只是四个数组的求和，在CPU上完成。ans是我需要的最终结果是的，但是你需要反向传播吗？@Ivan嗨，我不需要反向传播。谢谢你的回答，我会试试。一个问题是，它是按顺序运行还是同时运行？我认为不会，因为它不是异步代码。对于多GPU使用，您可能需要使用

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)