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"] =