首页 > 其他 > 详细

设置tensorflow的显存为动态使用

时间:2020-01-03 15:58:12      阅读:518      评论:0      收藏:0      [点我收藏+]

设置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的使用与分配

设置tensorflow的显存为动态使用

原文:https://www.cnblogs.com/yu212223/p/12145130.html

(0)
(0)
   
举报
评论 一句话评论(0
关于我们 - 联系我们 - 留言反馈 - 联系我们:wmxa8@hotmail.com
© 2014 bubuko.com 版权所有
打开技术之扣,分享程序人生!