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、自適應分配
會自適應分配顯存，不會將顯存全部分配導致資源浪費
|
1
2
3
|
config = tf.ConfigProto() config.gpu_options.allow_growth=Truesess = tf.Session(config=config) |
二、Keras
與tensorflow大差不差，就是將tf.Session配置轉置Keras配置
1、指定顯卡
代碼中加入
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
或者在運行代碼前，在終端
export CUDA_VISIBLE_DEVICES=0
2、為顯存分配使用比例
|
1
2
3
4
5
6
7
|
import tensorflow as tfimport keras.backend.tensorflow_backend as KTFconfig = tf.ConfigProto()config.gpu_options.per_process_gpu_memory_fraction = 0.333session = tf.Session(config=config)KTF.set_session(session) |
3、自適應分配
|
1
2
3
4
5
6
|
import keras.backend.tensorflow_backend as KTFconfig = tf.ConfigProto() config.gpu_options.allow_growth=Truesession = tf.Session(config=config)KTF.set_session(session) |
4、如有設置fit_generator
將多線程關閉
|
1
2
3
4
|
#可將use_multiprocessing=True#改為use_multiprocessing=False |
補充知識：Keras 自動分配顯存，不占用所有顯存
自動分配顯存，不占用所有顯存
|
1
2
3
4
5
6
7
8
9
|
import keras.backend.tensorflow_backend as KTFimport tensorflow as tfimport 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自適應使用顯存方式就是小編分享給大家的全部內容了，希望能給大家一個參考，也希望大家多多支持服務器之家。
原文鏈接：https://blog.csdn.net/l297969586/article/details/78905087
