Tensorflow與Keras自適應使用顯存方式

更新時間：2020年06月22日 10:19:44 作者：一呆飛仙

這篇文章主要介紹了Tensorflow與Keras自適應使用顯存方式，具有很好的參考價值，希望對大家有所幫助。一起跟隨小編過來看看吧

Tensorflow支持基于cuda內(nèi)核與cudnn的GPU加速，Keras出現(xiàn)較晚，為Tensorflow的高層框架，由于Keras使用的方便性與很好的延展性，之后更是作為Tensorflow的官方指定第三方支持開源框架。

但兩者在使用GPU時都有一個特點，就是默認為全占滿模式。在訓練的情況下，特別是分步訓練時會導致顯存溢出，導致程序崩潰。

可以使用自適應配置來調(diào)整顯存的使用情況。

一、Tensorflow

1、指定顯卡

代碼中加入

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

或者在運行代碼前，在終端

export CUDA_VISIBLE_DEVICES=0

2、為顯存分配使用比例

在建立tf.Session加入設置數(shù)據(jù)（顯存使用比例為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配置轉(zhuǎn)置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自適應使用顯存方式就是小編分享給大家的全部內(nèi)容了，希望能給大家一個參考，也希望大家多多支持腳本之家。

您可能感興趣的文章: