Skip to content

Instantly share code, notes, and snippets.

@s1998

s1998/reduce_mem_keras_tf.py

Last active May 9, 2021 20:04
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save s1998/be4c2cfc1629ad61e9312117b2988914 to your computer and use it in GitHub Desktop.
Save s1998/be4c2cfc1629ad61e9312117b2988914 to your computer and use it in GitHub Desktop.
Download ZIP
Restrict gpu memory in keras/tf
Raw
reduce_mem_keras_tf.py
# Set flexible GPU usage
import tensorflow as tf
import keras.backend as ktf
from tensorflow.python.keras import backend as ktf
def get_session(gpu_fraction=0.2):
gpu_options = tf.compat.v1.GPUOptions(
per_process_gpu_memory_fraction=gpu_fraction,
allow_growth=True)
return tf.compat.v1.Session(
config=tf.compat.v1.ConfigProto(gpu_options=gpu_options))
ktf.set_session(get_session())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment