TensorFlow 2.0 and Python 3 compatibility.

.gitignore

@@ -1,3 +1,7 @@
 venv/
 nasbench.egg-info/
 *.pyc
+
+# Data files
+nasbench_only108.tfrecord
+nasbench_full.tfrecord 

nasbench/api.py

@@ -99,7 +99,8 @@
 from nasbench.lib import model_metrics_pb2
 from nasbench.lib import model_spec as _model_spec
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
 
 # Bring ModelSpec to top-level for convenience. See lib/model_spec.py.
 ModelSpec = _model_spec.ModelSpec

nasbench/lib/base_ops.py

@@ -19,8 +19,8 @@
 from __future__ import print_function
 
 import abc
-
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
 
 # Currently, only channels_last is well supported.
 VALID_DATA_FORMATS = frozenset(['channels_last', 'channels_first'])

nasbench/lib/cifar.py

@@ -23,7 +23,8 @@
 from __future__ import print_function
 
 import functools
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
 
 WIDTH = 32
 HEIGHT = 32

nasbench/lib/evaluate.py

@@ -24,7 +24,8 @@
 from nasbench.lib import model_builder
 from nasbench.lib import training_time
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
 
 VALID_EXCEPTIONS = (
     tf.train.NanLossDuringTrainingError,  # NaN loss
@@ -122,7 +123,7 @@ def run(self):
         else:
           timing = training_time.limit(None)
 
-        evaluations = map(float, self.config['intermediate_evaluations'])
+        evaluations = list(map(float, self.config['intermediate_evaluations']))
         if not evaluations or evaluations[-1] != 1.0:
           evaluations.append(1.0)
         assert evaluations == sorted(evaluations)
@@ -197,8 +198,8 @@ def _evaluate_all(self, epochs, steps):
 
   def _compute_sample_metrics(self):
     """Computes the metrics on a fixed batch."""
-    sample_metrics = self.estimator.predict(
-        input_fn=self.input_sample.input_fn, yield_single_examples=False).next()
+    sample_metrics = next(self.estimator.predict(
+        input_fn=self.input_sample.input_fn, yield_single_examples=False))
 
     # Fix the extra batch dimension added by PREDICT
     for metric in sample_metrics:

nasbench/lib/model_builder.py

@@ -30,7 +30,8 @@
 from nasbench.lib import base_ops
 from nasbench.lib import training_time
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
 
 
 def build_model_fn(spec, config, num_train_images):
@@ -135,7 +136,7 @@ def model_fn(features, labels, mode, params):
       grads = tf.gradients(loss, all_params_tensors)
 
       param_gradient_norms = {}
-      for name, grad in zip(all_params_names, grads)[:-1]:
+      for name, grad in list(zip(all_params_names, grads))[:-1]:
         if grad is not None:
           param_gradient_norms[name] = (
               tf.expand_dims(tf.norm(grad, ord=2), 0))

nasbench/lib/training_time.py

@@ -19,7 +19,8 @@
 from __future__ import print_function
 
 import collections
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
 
 # Name of scope where to put all timing-related ops and variables.
 _SCOPE_NAME = 'timing'

nasbench/scripts/augment_model.py

@@ -22,7 +22,8 @@
 from nasbench.lib import evaluate
 from nasbench.lib import model_spec
 import numpy as np
-import tensorflow as tf     # Used for app, flags, logging
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
 
 tf.flags.DEFINE_string('model_dir', '', 'model directory')
 FLAGS = tf.flags.FLAGS

nasbench/scripts/generate_cifar10_tfrecords.py

@@ -35,7 +35,8 @@
 
 import tarfile
 from six.moves import cPickle as pickle
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
 
 CIFAR_FILENAME = 'cifar-10-python.tar.gz'
 CIFAR_DOWNLOAD_URL = 'https://www.cs.toronto.edu/~kriz/' + CIFAR_FILENAME

nasbench/scripts/generate_graphs.py

@@ -72,7 +72,9 @@
 
 from nasbench.lib import graph_util
 import numpy as np
-import tensorflow as tf   # For gfile
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
+
 
 flags.DEFINE_string('output_file', '/tmp/generated_graphs.json',
                     'Output file name.')

nasbench/scripts/run_evaluation.py

@@ -54,7 +54,8 @@
 from nasbench.lib import model_metrics_pb2
 from nasbench.lib import model_spec
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
 
 
 flags.DEFINE_string('models_file', '',

nasbench/tests/graph_util_test.py

@@ -21,7 +21,8 @@
 import random
 from nasbench.lib import graph_util
 import numpy as np
-import tensorflow as tf   # Used for tf.test
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
 
 
 class GraphUtilTest(tf.test.TestCase):

nasbench/tests/model_builder_test.py

@@ -20,7 +20,8 @@
 
 from nasbench.lib import model_builder
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
 
 
 class ModelBuilderTest(tf.test.TestCase):

nasbench/tests/model_spec_test.py

@@ -20,7 +20,8 @@
 
 from nasbench.lib import model_spec
 import numpy as np
-import tensorflow as tf   # Used only for tf.test
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
 
 
 class ModelSpecTest(tf.test.TestCase):

nasbench/tests/run_evaluation_test.py

@@ -24,8 +24,8 @@
 
 from absl.testing import flagsaver
 from nasbench.scripts import run_evaluation
-import tensorflow as tf
-
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
 
 
 class RunEvaluationTest(tf.test.TestCase):