浅谈多卡服务器下隐藏部分 GPU 和 TensorFlow 的显存使用设置

标题

浅谈多卡服务器下隐藏部分 GPU 和 TensorFlow 的显存使用设置

背景

在使用多卡服务器进行模型训练时，由于显存限制，可能会出现一部分 GPU 显存不足无法使用，或者需要提前将一些 GPU 预留出来供其他程序使用的情况。本文将提供一些方法来解决这种情况下的显存使用问题。

方案

1. 隐藏部分 GPU

在 Linux 系统下，我们可以使用 CUDA_VISIBLE_DEVICES 环境变量来控制 TensorFlow 使用哪些 GPU。该环境变量可以控制 TensorFlow 可见的 GPU 序号，从而可以隐藏一部分 GPU 不被 TensorFlow 使用。

例如，我们有一台 4 卡的服务器，但是我们只希望 TensorFlow 使用其中的 2 张卡，可以使用如下命令：

export CUDA_VISIBLE_DEVICES=0,1

这样就可以将只能够使用 0、1 号卡的显存留出来供其他程序使用。

2. TensorFlow 的显存使用设置

TensorFlow 中有一个重要的 GPU 设置，即 allow_growth。此设置控制 TensorFlow 是否在执行过程中逐渐申请 GPU 显存。在首次申请显存时，TensorFlow 会预留大量显存，可能会导致其他程序无法使用，也可能会导致一些 GPU 显存不足的情况。如果设置 allow_growth 为 True，则 TensorFlow 将会在需要时逐渐申请显存，从而可以避免显存预留过多的问题。

可以通过如下方式设置 allow_growth：

import tensorflow as tf

# 创建 Session
config = tf.ConfigProto()
# allow_growth 可避免显存占用过多
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)

示例

示例 1

我们有一台 4 卡的服务器，但是只有在使用 0、1 号卡进行模型训练时才能得到比较好的训练结果。我们希望将 2、3 号卡留出来供其他程序使用，同时设置 TensorFlow 的显存使用方式为 allow_growth。

我们可以在代码中加入如下代码：

import os
import tensorflow as tf

# 隐藏 2、3 号卡
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

# 创建 Session
config = tf.ConfigProto()
# allow_growth 可避免显存占用过多
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)

# 进行模型训练
...

这样，TensorFlow 只会使用 0、1 号卡进行模型训练，同时不会占用过多显存，从而能够使得其他程序可以正常运行。

示例 2

我们有一台 2 卡的服务器，但是其中一张卡的显存较小（仅为 2G），另一张卡显存较大（为 11G）。我们希望将模型训练放在显存较大的卡上，同时设置 TensorFlow 的显存使用方式为 allow_growth。

我们可以在代码中加入如下代码：

import os
import tensorflow as tf

# 隐藏显存较小的卡
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

# 创建 Session
config = tf.ConfigProto()
# allow_growth可避免显存占用过多
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)

# 进行模型训练
...

这样，TensorFlow 将会使用显存较大的卡进行模型训练，同时也不会占用过多显存，从而能够使得其他程序可以正常运行。

结论

通过使用上述方法，可以在使用多卡服务器进行模型训练时更好地控制显存的使用，避免显存被过分占用的问题，同时也能够更好地满足多程序共享服务器的需求。