diff --git a/tutorials/frontend/deploy_sparse.py b/tutorials/frontend/deploy_sparse.py index d3375c40fe72..f0af12b709e2 100644 --- a/tutorials/frontend/deploy_sparse.py +++ b/tutorials/frontend/deploy_sparse.py @@ -90,6 +90,20 @@ import scipy.sparse as sp +# Ask tensorflow to limit its GPU memory to what's actually needed +# instead of gobbling everything that's available. +# https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth +# This way this tutorial is a little more friendly to sphinx-gallery. +gpus = tf.config.list_physical_devices("GPU") +if gpus: + try: + for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + print("tensorflow will use experimental.set_memory_growth(True)") + except RuntimeError as e: + print("experimental.set_memory_growth option is not available: {}".format(e)) + + ############################################################################### # Configure Settings # ------------------ diff --git a/tutorials/frontend/from_tensorflow.py b/tutorials/frontend/from_tensorflow.py index fc87c07fb569..4563e245c0cf 100644 --- a/tutorials/frontend/from_tensorflow.py +++ b/tutorials/frontend/from_tensorflow.py @@ -36,6 +36,21 @@ # Tensorflow imports import tensorflow as tf + +# Ask tensorflow to limit its GPU memory to what's actually needed +# instead of gobbling everything that's available. +# https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth +# This way this tutorial is a little more friendly to sphinx-gallery. +gpus = tf.config.list_physical_devices("GPU") +if gpus: + try: + for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + print("tensorflow will use experimental.set_memory_growth(True)") + except RuntimeError as e: + print("experimental.set_memory_growth option is not available: {}".format(e)) + + try: tf_compat_v1 = tf.compat.v1 except ImportError: