使用keras框架训练神经网络时显卡内存使用很大很大的解决方案

Keras GPU内存设置

最新推荐文章于 2025-10-26 13:36:08 发布

原创最新推荐文章于 2025-10-26 13:36:08 发布 · 1.9k 阅读

0 ·

CC 4.0 BY-SA版权

keras 专栏收录该内容

6 篇文章

订阅专栏

本文详细介绍了如何在Keras中修改TensorFlow后端的源代码，以调整GPU内存使用比例，确保模型训练过程中GPU资源的有效利用。通过设定config.gpu_options.per_process_gpu_memory_fraction参数为0.3，可以限制每个进程使用的GPU内存，避免过度占用导致的训练失败。

修改源码：/usr/local/lib/python3.5/dist-packages/keras/backend/tensorflow_backend.py
else:
if _SESSION is None:
if not os.environ.get(‘OMP_NUM_THREADS’):
config = tf.ConfigProto(allow_soft_placement=True)
config.gpu_options.per_process_gpu_memory_fraction = 0.3 # zhex add
print(‘keras/backend/tensorflow_backend.py: config.gpu_options.per_process_gpu_memory_fraction = 0.3 #zhex’)
else:
num_thread = int(os.environ.get(‘OMP_NUM_THREADS’))
config = tf.ConfigProto(intra_op_parallelism_threads=num_thread, allow_soft_placement=True)
config.gpu_options.per_process_gpu_memory_fraction = 0.3 #yuzhijun add
print(‘config.gpu_options.per_process_gpu_memory_fraction = 0.3 #yuzhijun’)
_SESSION = tf.Session(config=config)
session = _SESSION