TensorFlow，Keras限制GPU显存

2024年5月26日 309次阅读 来源: tensorflow

　　运行TensorFlow程序会占用过多的显卡比例，多人共同使用GPU的时候，会造成后面的人无法运行程序。

  一、TensorFlow

　　1.预加载比例限制

tf_config = tensorflow.ConfigProto()  
tf_config.gpu_options.per_process_gpu_memory_fraction = 0.5 # 分配50%  
session = tensorflow.Session(config=tf_config)

　　2.自适应

tf_config = tensorflow.ConfigProto()  
tf_config.gpu_options.allow_growth = True  
session = tensorflow.Session(config=tf_config)

二、Keras

　　1.当使用Keras的情况下，当import keras时，框架会自动开启一个默认参数的Session。可以通过下面的代码主动创建一个使用了ConfigProto的Session，再注入到Keras的后端中去：

from keras.backend.tensorflow_backend import set_session  
config = tf.ConfigProto()  
config.gpu_options.allow_growth = True  
set_session(tf.Session(config=config))

   注意：第一行必须直接import这个函数，而非keras这个包，而且这句话必须在其他import keras之前，否则keras初始化之后，再替换一个session，原session也不会释放。

　　三、在Python程序中指定cpu

　　

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

 

    原文作者：tensorflow
    原文地址: https://www.cnblogs.com/youyou0/p/9547751.html
    本文转自网络文章，转载此文章仅为分享知识，如有侵权，请联系博主进行删除。
点赞