Last active
February 8, 2023 01:35
-
-
Save mexicantexan/b252e616b03e1b9b9f62f68874ffc974 to your computer and use it in GitHub Desktop.
Limit Keras GPU Memory with Tensorflow v2
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| """ | |
| For TensorFlow v2 and Keras you can limit gpu memory usage a couple of different ways, just add one of the following early on in your scripts | |
| """ | |
| import tensorflow as tf | |
| physical_devices = tf.config.experimental.list_physical_devices('GPU') | |
| tf.config.experimental.set_memory_growth(physical_devices[0], True) | |
| """ | |
| or | |
| """ | |
| import tensorflow as tf | |
| from keras.backend import set_session | |
| config = tf.compat.v1.ConfigProto() | |
| config.gpu_options.allow_growth = True | |
| config.gpu_options.visible_device_list = "0" # only the gpu 0 is allowed | |
| set_session(tf.compat.v1.Session(config=config)) | |
| """ | |
| For Tensorflow v1 | |
| """ | |
| import tensorflow as tf | |
| from keras.backend.tensorflow_backend import set_session | |
| config = tf.ConfigProto() | |
| config.gpu_options.allow_growth = True | |
| config.gpu_options.visible_device_list = "0" #only the gpu 0 is allowed | |
| set_session(tf.Session(config=config)) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment