TensorFlow 2.0 动态分配显存配置方法
发布时间：2026-10-16 |
13人浏览过
本文介绍在 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)) # 此处不同总结
以上为个人经验总结，希望能给大家一个参考，也希望大家多多支持本站。
本站声明
本文内容仅供参考，不代表本站观点。转载或引用请注明出处。如发现内容问题，请通过联系我们页面反馈。
