之前困扰笔者很久的一个 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。在终端中运行输出为：


MaxDirectMemorySize 配置_慢慢学TensorFlow

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


MaxDirectMemorySize 配置_应用程序_02

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

config.gpu_options.allow_growth=True

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


MaxDirectMemorySize 配置_github_03

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

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

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


MaxDirectMemorySize 配置_应用程序_04

上图中显示，/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 时，会经常用到这些参数。