tf_cnn_benchmark 显存问题

1 Overview

在测试 vGPU 的功能的时候，给容器分配了半张 GPU 卡，然后想用 Tensorflow Benchmark 测试一下，却发现半张 V100 32GB 显存从一开始就被占满了，但是 vGPU 只给 16 GB，这样发现 Tensorflow 无法正常运行。

2 Reason

因为 Tensorflow 默认会跑完整个显卡，为了可以测试顺利，所以想限制测试脚本 tf_cnn_benmark.py。看了一眼脚本，发现整个项目还依赖其他 python 文件，于是找一下，哪里可以把 Config 塞进去。

关于如何限制 GPU 显存的使用量，可以参考官方的记录。为了能正常使用办张卡，也就是 16 GB 的显存，我希望可以注入一个 Config 类似如下。

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.5)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

这样我的程序就只会用到半张卡的显存。

当然了，查看 tf_cnn_benchmark.py 文件，发现可以查看参数的。通过寻找选项，找到 --gpu_memory_frac_for_testing，这个参数实际就是设置 per_process_gpu_memory_fraction 的。

python tf_cnn_benchmark.py --help

于是最后我的执行脚本变成如下。

python
/examples/benchmarks/scripts/tf_cnn_benchmarks/tf_cnn_benchmarks.py
--num_gpus=1
--model=resnet50
--batch_size=32
--xla=True
--gpu_memory_frac_for_testing=0.5

3 Summary

per_process_gpu_memory_fraction 参数可以设置 GPU 内进程使用显存最大的比例。更多关于 Tensorflow 使用 GPU 的 tips，可以参考官方文档。