PyTorch 是一个非常优秀的深度学习框架，特别是在 GPU 上运行 PyTorch 代码能够加速深度网络的训练，但是，当服务器的 GPU 型号不匹配、性能差别较大、个人指定使用某个 GPU 时，需要进行特别的设置。本篇以 Ubuntu 18.04 为例，介绍如何在运行 PyTorch 代码时指定 GPU 训练模型。

代码中设置

当安装（深度学习的 GPU 环境的配置）了 GPU、cuda、cudnn 后，可以使用如下命令查看 GPU 信息：

1
2
3
$ nvidia-smi
或
$ nvidia-smi -L

结果大概如下：

1
2
3
GPU 0: Tesla V100-PCIE-32GB (UUID: GPU-xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxxx)
GPU 1: GeForce GTX TITAN X (UUID: GPU-xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxxx)
GPU 2: GeForce GTX TITAN X (UUID: GPU-xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxxx)

可见，服务器上安装了 3 块 GPU，但型号不同。

在编写 PyTorch 代码时，我们可以在主程序中增加如下代码，设置本程序运行只使用某些 GPU:

1
2
3
4
5
6
7
8
9
10
import os
# 先设置GPU显示顺序与 nvidia-smi 保持一致
# 此语句需在启动内核后，调用 torch 前设定
# torch 中GPU显示的顺序是按照 FASET_FIRST 顺序进行，有可能与 nvidia-smi 不一致
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"

# 只使用 GPU 0
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
# 只使用 GPU 1,2
os.environ["CUDA_VISIBLE_DEVICES"] = '1,2'

命令行设置

有时候克隆别人的代码，不便修改代码内容，此时，我们可以在命令行运行代码时同时设置指定的 GPU，方法如下：

1
2
3
4
5
# 只使用 GPU 0
CUDA_VISIBLE_DEVICES=0 python -m main

# 只使用 GPU 1,2
CUDA_VISIBLE_DEVICES=1,2 python -m main

环境变量设置

有时候实验室 GPU 资源紧张，安装的 GPU 都是分配到每个人使用的，那么个人只能使用固定的 GPU，此时，如果还向上面两种方法那样设置，会显得有些繁琐，需要每次运行时都要设置一次。通过个人用户下设置环境变量，可以限制每次运行都使用指定的 GPU，方法如下（假设使用的是 bash）：

1
2
3
cd 
echo "export CUDA_VISIBLE_DEVICES=0" >> .bashrc
source .bashrc

有些使用 pycharm 的人员可以参考下面的参考文献进行设置。

参考文献

  1. 选择要用的GPU: CUDA_VISIBLE_DEVICES
  2. os.environ[‘CUDA_VISIBLE_DEVICES‘]指定GPU无法生效解决方案
  3. 『Pytorch笔记』Pycharm中使用CUDA_VISIBLE_DEVICES=0！