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训练。

如图：
pytorch 指定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代码放在程序段最开始的部位。例如：
pytorch 指定GPU训练