Tensorflow设置显存自适应，显存比例-优快云博客

本文链接：https://blog.youkuaiyun.com/chanbo8205/article/details/85067855

本文详细介绍了如何在多GPU环境下管理GPU资源，包括通过环境变量和Python代码指定使用特定GPU的方法，以及如何在TensorFlow中精确配置GPU显存使用，既可按比例分配，也可按需动态增长。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1 终端执行程序时设置使用的GPU

如果电脑有多个GPU，tensorflow默认全部使用。如果想只使用部分GPU，可以设置CUDA_VISIBLE_DEVICES。在调用python程序时，可以使用（见第一个参考网址Franck Dernoncourt的回复）：

CUDA_VISIBLE_DEVICES=1 python my_script.py

Environment Variable Syntax      Results

CUDA_VISIBLE_DEVICES=1           Only device 1 will be seen
CUDA_VISIBLE_DEVICES=0,1         Devices 0 and 1 will be visible
CUDA_VISIBLE_DEVICES="0,1"       Same as above, quotation marks are optional
CUDA_VISIBLE_DEVICES=0,2,3       Devices 0, 2, 3 will be visible; device 1 is masked
CUDA_VISIBLE_DEVICES=""          No GPU will be visible

2 python代码中设置使用的GPU

如果要在python代码中设置使用的GPU（如使用pycharm进行调试时），可以使用下面的代码（见第二个参考网址中Yaroslav Bulatov的回复）：

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

3 设置tensorflow使用的显存大小

3.1 定量设置显存

默认tensorflow是使用GPU尽可能多的显存。可以通过下面的方式，来设置使用的GPU显存：

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

上面分配给tensorflow的GPU显存大小为：GPU实际显存*0.7。

可以按照需要，设置不同的值，来分配显存。

3.2 按需设置显存

上面的只能设置固定的大小。如果想按需分配，可以使用allow_growth参数（参考网址：http://blog.youkuaiyun.com/cq361106306/article/details/52950081）：

gpu_options = tf.GPUOptions(allow_growth=True)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

1. 按比例

1 config = tf.ConfigProto()
2 config.gpu_options.per_process_gpu_memory_fraction = 0.4
3 session = tf.Session(config=config, ...)

1     session_bundle_config.mutable_session_config()
2         ->mutable_gpu_options()
3         ->set_per_process_gpu_memory_fraction(0.4);

2.按需求增长

1 config = tf.ConfigProto()
2 config.gpu_options.allow_growth = True
3 session = tf.Session(config=config, ...)

1     session_bundle_config.mutable_session_config()
2         ->mutable_gpu_options()
3         ->set_allow_growth(true);