Merge pull request #39 from vlasenkoalexey/master

Making GraphDef Editor compatible with TensorFlow 2.x

README.md

@@ -126,5 +126,22 @@ $ pip install ./graph_def_editor
    `test.out` at the root of the project.
 
 
+## TensorFlow versions compartibility
+
+GraphDef Editor is fully supported for TensorFlow versions 1.14.x and 1.15.x.
+For TensorFlow 2.x some transforms might not work.
+
+To execute tests for specific TensorFlow version run following comand from the repository root:
+```sh
+docker run -v ${PWD}:/v -w /v tensorflow/tensorflow:<version>[-py3] bash -c "pip3 install -U pytest && pytest"
+```
+
+Pre 2.2.0 TensorFlow versions have -py3 suffix indicating that Python3 should be used.
+
+To execute specific test:
+```sh
+docker run -v ${PWD}:/v -w /v tensorflow/tensorflow:<version>[-py3] python -m tests.transform_test
+```
+
 
 

graph_def_editor/graph.py

@@ -22,7 +22,7 @@
 from distutils import dir_util
 import os
 from six import string_types
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import sys
 if sys.version >= '3':
   from typing import Tuple, Dict, FrozenSet, Iterable, Union, Set, Any

graph_def_editor/node.py

@@ -18,7 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import sys
 if sys.version >= '3':
   from typing import Tuple, List, Iterable, Any, AbstractSet

graph_def_editor/rewrite.py

@@ -24,7 +24,7 @@
 from __future__ import print_function
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import sys
 if sys.version >= '3':
   from typing import Tuple, Dict, Iterable, Union, Callable, Any

graph_def_editor/subgraph.py

@@ -24,7 +24,6 @@
 import six
 from six import iteritems
 from six import StringIO
-import tensorflow as tf
 
 from graph_def_editor import select, util
 from graph_def_editor import graph as gde_graph

graph_def_editor/tensor.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import sys
 if sys.version >= '3':
   from typing import AbstractSet

graph_def_editor/transform.py

@@ -25,7 +25,7 @@
 from six import iterkeys
 from six import string_types
 from six import StringIO
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import sys
 if sys.version >= '3':
   from typing import Iterable

graph_def_editor/util.py

@@ -28,7 +28,7 @@
 
 import numpy as np
 from six import iteritems, string_types
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from graph_def_editor import graph, node, tensor
 

tests/edit_test.py

@@ -19,7 +19,8 @@
 from __future__ import print_function
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_eager_execution()
 import unittest
 
 import graph_def_editor as gde

tests/graph_test.py

@@ -18,8 +18,8 @@
 """
 
 import unittest
-import tensorflow as tf
-import numpy as np
+import tensorflow.compat.v1 as tf
+tf.disable_eager_execution()
 import shutil
 import tempfile
 

tests/match_test.py

@@ -18,7 +18,9 @@
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_eager_execution()
+
 import unittest
 
 import graph_def_editor as gde

tests/reroute_test.py

@@ -19,7 +19,9 @@
 from __future__ import print_function
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_eager_execution()
+
 import unittest
 
 import graph_def_editor as gde

tests/rewrite_test.py

@@ -20,7 +20,9 @@
 import shutil
 import tempfile
 import unittest
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_eager_execution()
+
 import numpy as np
 
 import graph_def_editor as gde
@@ -759,3 +761,7 @@ def test_fold_batch_norms_up_fused_relu6(self):
     # Make sure the rewrite happened
     for n in g.nodes:
       self.assertNotEqual(n.op_type, "FusedBatchNorm")
+
+
+if __name__ == "__main__":
+  unittest.main()

tests/select_test.py

@@ -20,7 +20,9 @@
 
 import re
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_eager_execution()
+
 import unittest
 
 import graph_def_editor as gde

tests/subgraph_test.py

@@ -18,7 +18,9 @@
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_eager_execution()
+
 import unittest
 
 import graph_def_editor as gde

tests/transform_test.py

@@ -22,7 +22,9 @@
 import functools
 import numpy as np
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_eager_execution()
+
 import unittest
 
 import graph_def_editor as gde

tests/util_test.py

@@ -18,7 +18,9 @@
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_eager_execution()
+
 import unittest
 
 import graph_def_editor as gde
@@ -59,7 +61,7 @@ def test_unique_graph(self):
     a0, b0, a1, b1 = (g0["a"], g0["b"], g1["a"], g1["b"])
 
     print("g0['a'] returns {} (type {})".format(g0['a'], type(g0['a'])))
-    
+
     # Same graph, should be fine.
     self.assertIsNone(gde.util.check_graphs(a0, b0))
     # Two different graphs, should assert.
@@ -198,4 +200,4 @@ def test_identity(self):
 
 
 if __name__ == "__main__":
-  unittest.TestCase.main()
+  unittest.main()