(tensorflow) Fix issues when running with cuda

Solution from tensorflow/tensorflow#24496
tue-robotics · Feb 10, 2021 · bf5ebab · bf5ebab
1 parent 05af98a
commit bf5ebab
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/image_recognition_tensorflow/src/image_recognition_tensorflow/object_recognizer.py b/image_recognition_tensorflow/src/image_recognition_tensorflow/object_recognizer.py
@@ -32,7 +32,9 @@ def __init__(self, graph_path, labels_path, input_tensor, output_tensor):
             _ = tf.import_graph_def(graph_def, name='')
             tf.compat.v1.get_default_graph().finalize()  # Make the graph read-only, safe to use from any thread
 
-        self.session = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph())
+        config = tf.compat.v1.ConfigProto()
+        config.gpu_options.allow_growth = True
+        self.session = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph(), config=config)
 
         # This is only done to 'statically' check that the given tensor actually exists.
         # Not really statically but directly at startup, not when you need a first recognition

diff --git a/image_recognition_tensorflow/test/test_retrain_and_evaluate.py b/image_recognition_tensorflow/test/test_retrain_and_evaluate.py
@@ -20,5 +20,6 @@ def test_retrain_and_evaluate():
                                    os.path.join(assets_path, "verification"), "/tmp/result.csv")
     assert accuracy > 0.5
 
+
 if __name__ == "__main__":
     test_retrain_and_evaluate()