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

时间:2022-09-05 13:49:07

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

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

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

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

在第 211 行与 212 行之间增加一行,内容为:

config.gpu_options.allow_growth=True

保存为 modified.py,再次运行,输出如下:

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

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


TIPS:如何查看 GPU 显存占用情况?

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

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

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