GPU_ALLOW_GROWTH enabled in TF and TF v1

Author: lululxvi

deepxde/backend/tensorflow/__init__.py

@@ -1,9 +1,12 @@
 from __future__ import absolute_import
 
+import os
 from distutils.version import LooseVersion
 
 import tensorflow as tf
 
 
+os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"
+
 if LooseVersion(tf.__version__) < LooseVersion("2.2.0"):
     raise RuntimeError("DeepXDE requires tensorflow>=2.2.0.")

deepxde/backend/tensorflow_compat_v1/__init__.py

@@ -1,10 +1,13 @@
 from __future__ import absolute_import
 
+import os
 from distutils.version import LooseVersion
 
 import tensorflow as tf
 
 
+os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"
+
 if LooseVersion(tf.__version__) < LooseVersion("2.2.0"):
     raise RuntimeError("DeepXDE requires tensorflow>=2.2.0.")
 

deepxde/model.py

@@ -200,9 +200,7 @@ def predict(self, x, operator=None, callbacks=None):
     def _open_tfsession(self):
         if self.sess is not None:
             return
-        tfconfig = tf.ConfigProto()
-        tfconfig.gpu_options.allow_growth = True
-        self.sess = tf.Session(config=tfconfig)
+        self.sess = tf.Session()
         self.saver = tf.train.Saver(max_to_keep=None)
         self.train_state.set_tfsession(self.sess)