TensorFlow Session 中关于 GPU 的配置项解析 ——转自 慢慢学TensorFlow 微信公众号

之前困扰笔者很久的一个 TensorFlow 问题就是显存占用太高，跑一个最简单的 MNIST 例程居然就能把整个 GTX 1080 显存用满！导致其他需要使用 GPU 的程序运行时出现 Out of Memory 的错误。本文将介绍如何在创建 TensorFlow Session 对象时进行 GPU 相关设置避免上述问题。

TensorFlow 0.12 源码中提供了一个 alexnet_benchmark 例程，位于 https://github.com/tensorflow/tensorflow/blob/r0.12/tensorflow/models/image/alexnet/alexnet_benchmark.py，我们将该文件复制一份保存为 default.py。在终端中运行输出为：

运行时显存占用高达 7417 MiB！为了令其瘦身，我们需要修改 TensorFlow Session 创建部分的代码。

在第 211 行与 212 行之间增加一行，内容为：

config.gpu_options.allow_growth=True

保存为 modified.py，再次运行，输出如下：

运行时再次查看 GPU 显存占用变为 1739 MiB，剩下的显存可以让其他应用程序使用。

TIPS：如何查看 GPU 显存占用情况？

答：利用 NVIDIA 驱动中提供的 nvidia-smi 命令。如图所示：

上图中显示，/usr/bin/X 这个应用程序占用了 165 MiB 显存。此外也能看到 GTX 1080 显存总量为 8105 MiB，当前显存总占用为 398 MiB。更多信息请参考"nvidia-smi -h"。

虽然问题解决了，但好奇心驱使我对 config.gpu_options 一探究竟。

上面代码中使用的 config.gpu_options 实际上是用于配置 GPU 运行参数的 protocolbuffer 参数，为了搞明白其原理，需要首先看看 proto 源文件，位置：

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/protobuf/config.proto

摘录与 GPU 设备相关的内容如下：

message GPUOptions {

  // 取值范围 [0, 1]，表示为每个进程预先分配的 GPU 显存比例。 1 表示分配全部 GPU 显存， 0.5 表示分配可用 GPU 显存的一半

  double per_process_gpu_memory_fraction = 1;

  // 使用的 GPU 分配策略类型，"BFC"：最佳适配对齐算法，dlmalloc 简化的版本；默认为空字符串，使用系统选择的默认方法

  string allocator_type = 2;

  // 暂不关心

  int64 deferred_deletion_bytes = 3;

  // 允许显存增长。如果设置为 True，分配器不会预先分配一定量 GPU 显存，而是先分配一小块，必要时增加显存分配

  bool allow_growth = 4;

  // 用逗号分隔的一组 GPU ID，决定进程可见的 GPU 设备。

  string visible_device_list = 5;

  // 暂不关心

  int32 polling_active_delay_usecs = 6;

  // 暂不关心

  int32 polling_inactive_delay_msecs = 7;

};

读者可以在创建 TensorFlow Session 对象时尝试设置其他 config.gpu_options 参数，观察现象。在单机多卡运行 TensorFlow 时，会经常用到这些参数。