pytorch 指定GPU训练
指定GPU代码
上述使用中,代码3最为常用,keras也可以通过这种方式来调用GPU训练神经网络。
# 代码1:
torch.cuda.set_device(1)
# 代码2:
device = torch.device("cuda:1")
# 代码3:(官方推荐使用)
os.environ["CUDA_VISIBLE_DEVICES"] = '1'
(如果你想同时调用两块GPU的话)
os.environ["CUDA_VISIBLE_DEVICES"] = '1,2'
使用过程中遇到的问题
问题:在使用上述代码过程中,程序仍然使用默认GPU训练。
如图:
之后,程序报错:
RuntimeError: Expected tensor for argument #1 'input' to have the same device as tensor for argument #2 'weight'; but device 1 does not equal 0 (while checking arguments for cudnn_convolution)
那么,你可能需要将指定GPU代码放在程序段最开始的部位。例如: