This is an automated email from the ASF dual-hosted git repository.
junrushao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new cfa498c Make from_tensorflow.py more GPU memory friendly. (#8763)
cfa498c is described below
commit cfa498c0376622afe4e0f7344f0104dc97d7e876
Author: Mark Shields <[email protected]>
AuthorDate: Mon Aug 16 22:44:30 2021 -0700
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 d3375c4..f0af12b 100644
--- a/tutorials/frontend/deploy_sparse.py
+++ b/tutorials/frontend/deploy_sparse.py
@@ -90,6 +90,20 @@ from tensorflow.python.framework.convert_to_constants import
(
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 fc87c07..4563e24 100644
--- a/tutorials/frontend/from_tensorflow.py
+++ b/tutorials/frontend/from_tensorflow.py
@@ -36,6 +36,21 @@ import os.path
# 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:
