tf_cnn_benchmark 显存问题

2020-08-05 23:53:50 浏览数 (1)

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 类似如下。

代码语言:javascript复制
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 的。

代码语言:javascript复制
python tf_cnn_benchmark.py --help

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

代码语言:javascript复制
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,可以参考官方文档。

0 人点赞