python

关注公众号 jb51net

关闭
首页 > 脚本专栏 > python > Tensorflow Keras自适应显存

Tensorflow与Keras自适应使用显存方式

作者:一呆飞仙

这篇文章主要介绍了Tensorflow与Keras自适应使用显存方式,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧

Tensorflow支持基于cuda内核与cudnn的GPU加速,Keras出现较晚,为Tensorflow的高层框架,由于Keras使用的方便性与很好的延展性,之后更是作为Tensorflow的官方指定第三方支持开源框架。

但两者在使用GPU时都有一个特点,就是默认为全占满模式。在训练的情况下,特别是分步训练时会导致显存溢出,导致程序崩溃。

可以使用自适应配置来调整显存的使用情况。

一、Tensorflow

1、指定显卡

代码中加入

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

或者在运行代码前,在终端

export CUDA_VISIBLE_DEVICES=0

2、为显存分配使用比例

在建立tf.Session加入设置数据(显存使用比例为1/3),但有时你虽然设置了使用上限,在程序需要更高显存时还是会越过该限制

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

3、自适应分配

会自适应分配显存,不会将显存全部分配导致资源浪费

config = tf.ConfigProto() 
config.gpu_options.allow_growth=True 
sess = tf.Session(config=config) 

二、Keras

与tensorflow大差不差,就是将tf.Session配置转置Keras配置

1、指定显卡

代码中加入

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

或者在运行代码前,在终端

export CUDA_VISIBLE_DEVICES=0

2、为显存分配使用比例

import tensorflow as tf
import keras.backend.tensorflow_backend as KTF

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.333
session = tf.Session(config=config)
KTF.set_session(session)

3、自适应分配

import keras.backend.tensorflow_backend as KTF

config = tf.ConfigProto() 
config.gpu_options.allow_growth=True 
session = tf.Session(config=config)
KTF.set_session(session)

4、如有设置fit_generator

将多线程关闭

#可将
use_multiprocessing=True
#改为
use_multiprocessing=False

补充知识:Keras 自动分配显存,不占用所有显存

自动分配显存,不占用所有显存

import keras.backend.tensorflow_backend as KTF
import tensorflow as tf
import os
 
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
config = tf.ConfigProto()
config.gpu_options.allow_growth=True #不全部占满显存, 按需分配
sess = tf.Session(config=config)
KTF.set_session(sess)

以上这篇Tensorflow与Keras自适应使用显存方式就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持脚本之家。

您可能感兴趣的文章:
阅读全文