Pytorch指定GPU进行训练的方法

Pytorch指定GPU进行训练的方法

  在用Pytorch深度学习框架进行训练模型的时候，默认使用服务器的第“0”号GPU训练。当“0”号GPU被占用时，不能正常进行训练，通常会报错“CUDA out of memory”，这时需要选用其余GPU进行训练。用os模块即可解决：

  如上图所示，此时第“0”号GPU已被占用，则需要在用于训练的.py文件中加入如下的代码(一般在import之后加入)：

import os

#"1, 2"表示训练的时候选用两块GPU，优先选用"1"号GPU
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1, 2" 

  如果使用一块GPU进行训练，则代码相应为：

import os

#"1"表示训练的时候选用"1"号GPU
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1" 

  需要注意的是，如果训练的.py文件有引入自己编写的其它.py文件，如utilities.py、models.py等这些文件里面定义了一些训练时所用到的模型和方法，那么也需要在这些.py文件中加入上述代码，这样才可正确地指定GPU训练模型。