From cfa498c0376622afe4e0f7344f0104dc97d7e876 Mon Sep 17 00:00:00 2001
From: Mark Shields <87091372+mbs-octoml@users.noreply.github.com>
Date: Mon, 16 Aug 2021 22:44:30 -0700
Subject: [PATCH] Make from_tensorflow.py more GPU memory friendly. (#8763)

* Make from_tensorflow.py more GPU memory friendly.

Sphinx-gallery runs everything in a single process. There
doesn't appear to be any easy way to force Tensorflow to
return memory other than terminating the process. This at
least gives us a little more wiggle room.

* Also deploy_sparse.py. Should probably also be done to tensorflow.rst.
---
 tutorials/frontend/deploy_sparse.py   | 14 ++++++++++++++
 tutorials/frontend/from_tensorflow.py | 15 +++++++++++++++
 2 files changed, 29 insertions(+)

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: