Tensorflow與Keras自適應使用顯存

Tensorflow支持基於cuda內核與cudnn的GPU加速，Keras出現較晚，為Tensorflow的高層框架，由於Keras使用的方便性與很好的延展性，之后更是作為Tensorflow的官方指定第三方支持開源框架。但兩者在使用GPU時都有一個特點，就是默認為全占滿模式。在訓練的情況下，特別是分步訓練時會導致顯存溢出，導致程序崩潰。可以使用自適應配置來調整顯存的使用情況。

一、Tensorflow
1、指定顯卡
代碼中加入

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
1
2
或者在運行代碼前，在終端

export CUDA_VISIBLE_DEVICES=0
1
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))
1
2
3、自適應分配
會自適應分配顯存，不會將顯存全部分配導致資源浪費

config = tf.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.Session(config=config)
1
2
3
二、Keras
與tensorflow大差不差，就是將tf.Session配置轉置Keras配置
1、指定顯卡
代碼中加入

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
1
2
或者在運行代碼前，在終端

export CUDA_VISIBLE_DEVICES=0
1
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)
1
2
3
4
5
6
7
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)
1
2
3
4
5
6
4、如有設置fit_generator
將多線程關閉

#可將
use_multiprocessing=True
#改為
use_multiprocessing=False
————————————————
版權聲明：本文為CSDN博主「一呆飛仙」的原創文章，遵循 CC 4.0 BY-SA 版權協議，轉載請附上原文出處鏈接及本聲明。
原文鏈接：https://blog.csdn.net/l297969586/article/details/78905087