假如你有多个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])