fix broken test cases

Author: fsx950223

efficientdet/tf2/efficientdet_keras.py

@@ -451,8 +451,8 @@ def conv2d_layer(cls, separable_conv, data_format):
           tf.keras.layers.SeparableConv2D,
           depth_multiplier=1,
           data_format=data_format,
-          pointwise_initializer=tf.initializers.variance_scaling(),
-          depthwise_initializer=tf.initializers.variance_scaling())
+          pointwise_initializer='variance_scaling',
+          depthwise_initializer='variance_scaling')
     else:
       conv2d_layer = functools.partial(
           tf.keras.layers.Conv2D,
@@ -537,8 +537,8 @@ def __init__(self,
             tf.keras.layers.SeparableConv2D(
                 filters=self.num_filters,
                 depth_multiplier=1,
-                pointwise_initializer=tf.initializers.variance_scaling(),
-                depthwise_initializer=tf.initializers.variance_scaling(),
+                pointwise_initializer='variance_scaling',
+                depthwise_initializer='variance_scaling',
                 data_format=self.data_format,
                 kernel_size=3,
                 activation=None,
@@ -612,8 +612,8 @@ def boxes_layer(cls, separable_conv, num_anchors, data_format, name):
       return tf.keras.layers.SeparableConv2D(
           filters=4 * num_anchors,
           depth_multiplier=1,
-          pointwise_initializer=tf.initializers.variance_scaling(),
-          depthwise_initializer=tf.initializers.variance_scaling(),
+          pointwise_initializer='variance_scaling',
+          depthwise_initializer='variance_scaling',
           data_format=data_format,
           kernel_size=3,
           activation=None,

efficientdet/tf2/efficientdet_keras_test.py

@@ -17,6 +17,7 @@
 import tempfile
 from absl import logging
 import tensorflow.compat.v1 as tf
+import tensorflow.compat.v2 as tf2
 import efficientdet_arch as legacy_arch
 import hparams_config
 from tf2 import efficientdet_keras
@@ -38,26 +39,37 @@ def test_model_output(self):
     inputs_shape = [1, 512, 512, 3]
     config = hparams_config.get_efficientdet_config('efficientdet-d0')
     config.heads = ['object_detection', 'segmentation']
+    tf2.keras.utils.set_random_seed(SEED)
     with tf.Session(graph=tf.Graph()) as sess:
       feats = tf.ones(inputs_shape)
-      tf.random.set_random_seed(SEED)
       model = efficientdet_keras.EfficientDetNet(config=config)
       outputs = model(feats, True)
       sess.run(tf.global_variables_initializer())
       keras_class_out, keras_box_out, _ = sess.run(outputs)
       grads = tf.nest.map_structure(lambda output: tf.gradients(output, feats),
                                     outputs)
       keras_class_grads, keras_box_grads, _ = sess.run(grads)
+      vars = list(filter(
+          lambda var: not var.name.startswith('segmentation'),
+          tf.global_variables()))
+      vars.sort(key=lambda var: var.name)
+      keras_vars_names = [var.name for var in vars]
+      keras_vars_values = sess.run(vars)
+
     with tf.Session(graph=tf.Graph()) as sess:
       feats = tf.ones(inputs_shape)
-      tf.random.set_random_seed(SEED)
       outputs = legacy_arch.efficientdet(feats, config=config)
-      sess.run(tf.global_variables_initializer())
+      vars = tf.global_variables()
+      vars.sort(key=lambda var: var.name)
+      legacy_vars_names = [var.name for var in vars]
+      sess.run([var.assign(val) for val, var in zip(keras_vars_values, vars)])
       legacy_class_out, legacy_box_out = sess.run(outputs)
       grads = tf.nest.map_structure(lambda output: tf.gradients(output, feats),
                                     outputs)
       legacy_class_grads, legacy_box_grads = sess.run(grads)
 
+    self.assertAllEqual(keras_vars_names, legacy_vars_names)
+
     for i in range(3, 8):
       self.assertAllEqual(
           keras_class_out[i - 3], legacy_class_out[i])
@@ -76,7 +88,7 @@ def test_eager_output(self):
 
     with tf.Session(graph=tf.Graph()) as sess:
       feats = tf.ones(inputs_shape)
-      tf.random.set_random_seed(SEED)
+      tf2.keras.utils.set_random_seed(SEED)
       model = efficientdet_keras.EfficientDetNet(config=config)
       outputs = model(feats, True)
       grads = tf.nest.map_structure(lambda output: tf.gradients(output, feats),
@@ -120,7 +132,7 @@ def test_build_feature_network(self):
           tf.ones([1, 32, 32, 112]),  # level 4
           tf.ones([1, 16, 16, 320]),  # level 5
       ]
-      tf.random.set_random_seed(SEED)
+      tf2.keras.utils.set_random_seed(SEED)
       fpn_cell = efficientdet_keras.FPNCells(config)
       new_feats1 = fpn_cell(inputs, True)
       sess.run(tf.global_variables_initializer())
@@ -137,7 +149,7 @@ def test_build_feature_network(self):
           4: tf.ones([1, 32, 32, 112]),
           5: tf.ones([1, 16, 16, 320])
       }
-      tf.random.set_random_seed(SEED)
+      tf2.keras.utils.set_random_seed(SEED)
       new_feats2 = legacy_arch.build_feature_network(inputs, config)
       sess.run(tf.global_variables_initializer())
       legacy_feats = sess.run(new_feats2)
@@ -185,7 +197,7 @@ def test_resample_feature_map(self):
         for strategy in ['tpu', '']:
           with self.subTest(
               apply_bn=apply_bn, training=training, strategy=strategy):
-            tf.random.set_random_seed(SEED)
+            tf2.keras.utils.set_random_seed(SEED)
             expect_result = legacy_arch.resample_feature_map(
                 feat,
                 name='resample_p0',
@@ -195,7 +207,7 @@ def test_resample_feature_map(self):
                 apply_bn=apply_bn,
                 is_training=training,
                 strategy=strategy)
-            tf.random.set_random_seed(SEED)
+            tf2.keras.utils.set_random_seed(SEED)
             resample_layer = efficientdet_keras.ResampleFeatureMap(
                 name='resample_p0',
                 feat_level=0,

efficientdet/tf2/infer_lib.py

@@ -214,7 +214,7 @@ def predict(self, image_arrays):
     Returns:
       Model outputs.
     """
-    raise NotImplemented
+    raise NotImplementedError
 
   def _preprocess(self, image_arrays):
 
@@ -249,7 +249,7 @@ def serve(self, image_arrays):
     Returns:
       A list of detections.
     """
-    raise NotImplemented
+    raise NotImplementedError
 
   def _get_model_and_spec(self, tflite=None):
     """Get model instance and export spec."""

efficientdet/tf2/infer_lib_test.py

@@ -90,7 +90,7 @@ def test_infer_lib(self):
     images = tf.ones((1, 512, 512, 3))
     boxes, scores, classes, valid_lens = driver.serve(images)
     self.assertEqual(tf.reduce_mean(boxes), 163.09)
-    self.assertEqual(tf.reduce_mean(scores), 0.01000005)
+    self.assertEqual(tf.reduce_mean(scores), 0.01)
     self.assertEqual(tf.reduce_mean(classes), 1)
     self.assertEqual(tf.reduce_mean(valid_lens), 100)
     self.assertEqual(boxes.shape, (1, 100, 4))
@@ -107,7 +107,7 @@ def test_infer_lib_without_ema(self):
     images = tf.ones((1, 512, 512, 3))
     boxes, scores, classes, valid_lens = driver.serve(images)
     self.assertEqual(tf.reduce_mean(boxes), 163.09)
-    self.assertEqual(tf.reduce_mean(scores), 0.01000005)
+    self.assertEqual(tf.reduce_mean(scores), 0.01)
     self.assertEqual(tf.reduce_mean(classes), 1)
     self.assertEqual(tf.reduce_mean(valid_lens), 100)
     self.assertEqual(boxes.shape, (1, 100, 4))
@@ -134,7 +134,7 @@ def test_infer_lib_mixed_precision(self):
     policy = tf.keras.mixed_precision.global_policy()
     if policy.name == 'float32':
       self.assertEqual(tf.reduce_mean(boxes), 163.09)
-      self.assertEqual(tf.reduce_mean(scores), 0.01000005)
+      self.assertEqual(tf.reduce_mean(scores), 0.01)
       self.assertEqual(tf.reduce_mean(classes), 1)
       self.assertEqual(tf.reduce_mean(valid_lens), 100)
     elif policy.name == 'float16':

efficientdet/tf2/train_lib.py

@@ -661,7 +661,10 @@ def train_step(self, data):
         scaled_loss = total_loss
         optimizer = self.optimizer
     loss_vals['loss'] = total_loss
-    loss_vals['learning_rate'] = optimizer.learning_rate(optimizer.iterations)
+    if callable(optimizer.learning_rate):
+      loss_vals['learning_rate'] = optimizer.learning_rate(optimizer.iterations)
+    else:
+      loss_vals['learning_rate'] = optimizer.learning_rate
     trainable_vars = self._freeze_vars()
     scaled_gradients = tape.gradient(scaled_loss, trainable_vars)
     if isinstance(self.optimizer,

efficientdet/tf2/train_lib_test.py

@@ -196,8 +196,8 @@ def test_fit(self):
     # skip gnorm test because it is flaky.
 
   def test_recompute_grad(self):
-    tf.config.run_functions_eagerly(True)
     _, x, labels, model = self._build_model(False)
+    weights = model.get_weights()
     with tf.GradientTape() as tape:
       loss_vals = {}
       cls_outputs, box_outputs, _ = model(x, training=True)
@@ -206,6 +206,7 @@ def test_recompute_grad(self):
       grads1 = tape.gradient(det_loss, model.trainable_variables)
 
     _, x, labels, model = self._build_model(True)
+    model.set_weights(weights)
     with tf.GradientTape() as tape:
       loss_vals = {}
       cls_outputs2, box_outputs2, _ = model(x, training=True)