TensorFlow 2.0 动态分配显存配置方法

Finn

发布时间：2026-10-16| 12人浏览过

原创

本文介绍在 TensorFlow 2.0 及基于 Keras 的环境中，如何配置 GPU 动态分配显存。通过设置 ConfigProto 的 allow_growth 或 per_process_gpu_memory_fraction 参数，避免程序一次性占用所有 GPU 资源，从而支持多用户或多模型共享服务器 GPU。

TensorFlow 2.0 之后动态分配显存

import tensorflow as tf
 
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.compat.v1.Session(config = config)
tf.compat.v1.keras.backend.set_session(sess)

keras和tensorflow动态分配显存

在使用 TensorFlow 或基于 TensorFlow 的 Keras 时，默认行为往往是将 GPU 中的所有显存预先读取并占用，随后再逐步使用。

如果服务器仅有一张 GPU，这种默认机制会导致其他用户无法使用显卡资源。因此，这里介绍一种动态分配显存的方法：

在 TensorFlow 环境下，可以通过添加以下代码来实现：

config = tf.ConfigProto()
config.gpu_options.allow_growth = True   #设置动态分配显存
session = tf.Session(config=config, ...)

或者采用如下方式：

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.3   #只一次调用30%的显存
session = tf.Session(config=config, ...)

对于 Keras，若使用 TensorFlow 作为后端，只需额外增加一句配置即可：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.3
set_session(tf.Session(config=config)) # 此处不同