如何解决Torch,如何为不同的数据集使用多个 GPU
假设我有 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 个 ans 移回 cpu 然后计算总和
ans = ans0 + ans1 + ans2 + ans3
解决方法
假设您只需要 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)
版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。