设置tensorflow的显存为动态使用


设置tensorflow的显存为动态使用

默认情况下,TensorFlow 将使用几乎所有可用的显存,以避免内存碎片化所带来的性能损失,但这样不能在一台机器上运行多个程序

tensorflow 1.x

from tensorflow.compat.v1 import GPUOptions
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import Session
from keras.backend.tensorflow_backend import set_session

gpu_options = GPUOptions(allow_growth=True)
set_session(Session(config=ConfigProto(gpu_options=gpu_options)))

详见tensorflow入门笔记1:指定GPU及分配显存

tensorflow 2.0

import tensorflow as tf

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

详见Tensorflow 2.0 GPU的使用与分配


免责声明!

本站转载的文章为个人学习借鉴使用,本站对版权不负任何法律责任。如果侵犯了您的隐私权益,请联系本站邮箱yoyou2525@163.com删除。



 
粤ICP备18138465号  © 2018-2025 CODEPRJ.COM