search

Link-Li / Balanced-DataParallel

这里是改进了pytorch的DataParallel, 用来平衡第一个GPU的显存使用量
230 stars 51 forks source link

BalancedDataParallel 如何指定显卡 #12

Closed QUTGXX closed 3 years ago

QUTGXX commented 3 years ago

您好 我想用4快卡的机器训练，但只想其中2块卡运行此程序，请问如何设定呢？ 我已尝试修改self.device_ids， 但其仍使用所有显卡训练。可以帮助我解决此问题吗？

Link-Li commented 3 years ago

用CUDA_VISIBLE_DEVICES关键字设置你要用的显卡

QUTGXX commented 3 years ago

CUDA_VISIBLE_DEVICES

用了，但还是使用4块卡。

Link-Li commented 3 years ago

代码里面使用os.environ['CUDA_VISIBLE_DEVICES']进行设置你要使用的显卡

QUTGXX commented 3 years ago

代码里面使用os.environ['CUDA_VISIBLE_DEVICES']进行设置你要使用的显卡

对 我就是这么设置的。 但还是用了4块卡。 部分代码如下： os.environ['CUDA_VISIBLE_DEVICES'] = "0,2"

netD.to(gpu_ids[0]) netD = BalancedDataParallel(16,netD, dim=0)

在class BalancedDataParallel(DataParallel)中修改如下： device_ids = [0,2] inputs, kwargs = self.scatter(inputs, kwargs, device_ids)

Link-Li commented 3 years ago

那就不知道了，建议谷歌找答案

QUTGXX commented 3 years ago

那就不知道了，建议谷歌找答案

你好，我通过'CUDA_VISIBLE_DEVICES=0,2 python ....' 解决了此问题。 目前设置0号卡batchsize为8 总batchsize为32. 但训练时间比之前直接使用torch.nn.DataParallel (batchsize 24)慢25%。 您知道如何解决吗？