pytorch多GPU的正确使用方式

假如你有多个GPU(>2),此处假设你有8块,并且前4块已经使用了,那么怎么才能从第5块开始用呢?
常见错误:

    raise AssertionError("Invalid device id")
AssertionError: Invalid device id
THCudaCheck FAIL file=..\torch\csrc\cuda\Module.cpp line=59 error=101 : invalid
device ordinal
Traceback (most recent call last):
  File "train2.py", line 26, in <module>
    torch.cuda.set_device(4)
  File "C:\ProgramData\Anaconda3\lib\site-packages\torch\cuda\__init__.py", line
 281, in set_device
    torch._C._cuda_setDevice(device)
RuntimeError: cuda runtime error (101) : invalid device ordinal at ..\torch\csrc
\cuda\Module.cpp:59

解决方法:
我们知道GPU的编号是从0开始的,这8块的编号就是0,1,2,3,4,5,6,7
现在要从4开始使用,只用4和5
1.在你的训练脚本import 结束的时候放下面两句:

torch.cuda.set_device(4)  #指定主显卡,不指定不行,默认是0,但是0被占用了
os.environ["CUDA_VISIBLE_DEVICES"] = "4,5"

2.脚本里面需要放:

device = torch.device("cuda:4" if torch.cuda.is_available() else "cpu")

3.还需要改:

self.net = torch.nn.DataParallel(self.net, device_ids=[4,5])

猜你喜欢

转载自blog.csdn.net/qq_20373723/article/details/109252699