Updated silence tensorflow for tensorflow 2.0

LucaCappelletti94 · Mar 2, 2020 · 6a51f91 · 6a51f91
1 parent da9c303
commit 6a51f91
Show file tree

Hide file tree

Showing 10 changed files with 20 additions and 135 deletions.
diff --git a/.gitignore b/.gitignore
@@ -12,5 +12,4 @@
 
 .notipy.json
 \.DS_Store
-
-
+*.egg-info
diff --git a/silence_tensorflow.egg-info/PKG-INFO b/silence_tensorflow.egg-info/PKG-INFO
diff --git a/silence_tensorflow.egg-info/SOURCES.txt b/silence_tensorflow.egg-info/SOURCES.txt
diff --git a/silence_tensorflow.egg-info/dependency_links.txt b/silence_tensorflow.egg-info/dependency_links.txt
diff --git a/silence_tensorflow.egg-info/requires.txt b/silence_tensorflow.egg-info/requires.txt
diff --git a/silence_tensorflow.egg-info/top_level.txt b/silence_tensorflow.egg-info/top_level.txt
diff --git a/silence_tensorflow/__init__.py b/silence_tensorflow/__init__.py
@@ -1,8 +1,5 @@
-import os
-os.environ["KMP_AFFINITY"] = "noverbose"
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
-import warnings
-warnings.simplefilter(action='ignore', category=FutureWarning)
-warnings.simplefilter(action='ignore', category=DeprecationWarning)
-import tensorflow as tf
-tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
+from .silence_tensorflow import silence_tensorflow
+
+__all__ = [
+    "silence_tensorflow"
+]
diff --git a/silence_tensorflow/silence_tensorflow.py b/silence_tensorflow/silence_tensorflow.py
@@ -0,0 +1,12 @@
+import os
+import logging
+
+
+def silence_tensorflow():
+    """Silence every warning of notice from tensorflow."""
+    logging.getLogger('tensorflow').setLevel(logging.ERROR)
+    os.environ["KMP_AFFINITY"] = "noverbose"
+    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
+    import tensorflow as tf
+    tf.get_logger().setLevel('ERROR')
+    tf.autograph.set_verbosity(3)
diff --git a/tests/test_import.py b/tests/test_import.py
diff --git a/tests/test_silence_tensorflow.py b/tests/test_silence_tensorflow.py
@@ -1,7 +1,5 @@
-import silence_tensorflow
-import os
+from silence_tensorflow import silence_tensorflow
 
 def test_silence_tensorflow():
     """Check that everything runs."""
-    assert os.environ["KMP_AFFINITY"] == "noverbose"
-    assert os.environ["TF_CPP_MIN_LOG_LEVEL"] == "2"
+    silence_tensorflow()