TensorFlow GPU 设置

TensorFlow 的默认配置是使用所有可见的 GPU 的所有显存,这在某些程度上带来了不便。

使用指定的 GPU

命令行执行时指定

使用 CUDA_VISIBLE_DEVICES 来指定脚本的可见 GPU

1
CUDA_VISIBLE_DEVICES=1 python example.py

可以设置可见多个 GPU:

1
2
CUDA_VISIBLE_DEVICES=0,1,2
CUDA_VISIBLE_DEVICES=""

在代码中指定

可以使用 os 来设置当前脚本的可见环境

1
2
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

相比第一种方式,缺少了一些灵活性。

tf.ConfigProto

1
2
3
4
5
tf.ConfigProto(
allow_soft_placement=True,
gpu_options=tf.GPUOptions(),
log_device_placement=True
)

tf.ConfigProto 用来对 session 进行参数配置,常用参数:

  1. allow_soft_placement=True:自动选择可运行的设备。如果程序中指定了设备,在不可用时会抛出异常,这里为了防止等待或者异常,允许 TensorFlow 自己选择合适的设备;
  2. gpu_options=tf.GPUOptions():设置 GPU 的显存占用等,见后文;
  3. log_device_placement=True:记录设备指派情况。打印出所有 operations 和 Tensor 在哪个设备上运行。

tf.GPUOptions

1
2
3
4
tf.GPUOptions(
allow_growth,
per_process_gpu_memory_fraction,
visible_device_list)

tf.GPUOptions 作为 tf.ConfigProto 的一个参数,用来设置脚本运行的 GPU 资源,常用参数:

  1. allow_growth=True:动态申请显存。开始分配一小部分显存,随着需要而增加;
  2. per_process_gpu_memory_fraction=0.7:预分配的 GPU 显存占用最大比例。

常用设置

1
2
3
4
5
gpu_option = tf.GPUOptions(allow_growth=True, per_process_gpu_memory_fraction=...)
config = tf.ConfigProto(allow_soft_placement=True, gpu_options=gpu_option, log_device_placement=True)
with tf.Session(config=config) as sess:
sess.run(...)
...

参考

https://www.jianshu.com/p/b9a442bcfd2e
https://blog.csdn.net/c_chuxin/article/details/84990176