Merged commit includes the following changes: (#7354)

261171038  by gjn<gjn@google.com>:

    Remove weight_decay_rate 0 early exit check

    Removing this code path should be fine since this was actually not doing
    what it meant to do. Since weight_decay_rate is actually a tensor, the
    equality check was only looking at the id of the object and comparing to
    0. This should never be true. Evaluating a tensor is also not what we
    want to do at this point of the code. Thus it should be fine to simply
    remove this code.

--
261169862  by haoyuzhang<haoyuzhang@google.com>:

    Internal change

261153520  by haoyuzhang<haoyuzhang@google.com>:

    Internal change

261140302  by hongkuny<hongkuny@google.com>:

    Clean up

--

PiperOrigin-RevId: 261171038

official/bert/modeling.py

@@ -799,10 +799,6 @@ class Transformer(tf.keras.layers.Layer):
               name=("layer_%d" % i)))
     super(Transformer, self).build(unused_input_shapes)
 
-    # Workaround for Keras bug where layers aren't tracked properly.
-    for i in range(len(self.layers)):
-      self.__setattr__("layer%d" % i, self.layers[i])
-
   def __call__(self, input_tensor, attention_mask=None):
     inputs = pack_inputs([input_tensor, attention_mask])
     return super(Transformer, self).__call__(inputs=inputs)

official/bert/optimization.py

@@ -169,8 +169,6 @@ class AdamWeightDecay(tf.keras.optimizers.Adam):
 
   def _do_use_weight_decay(self, param_name):
     """Whether to use L2 weight decay for `param_name`."""
-    if self.weight_decay_rate == 0:
-      return False
     if self._exclude_from_weight_decay:
       for r in self._exclude_from_weight_decay:
         if re.search(r, param_name) is not None:

official/resnet/keras/keras_imagenet_main.py

@@ -270,7 +270,8 @@ def define_imagenet_keras_flags():
 def main(_):
   model_helpers.apply_clean(flags.FLAGS)
   with logger.benchmark_context(flags.FLAGS):
-    run(flags.FLAGS)
+    stats = run(flags.FLAGS)
+  tf.compat.v1.logging.info('Run stats:\n%s' % stats)
 
 
 if __name__ == '__main__':

official/transformer/v2/beam_search.py

@@ -49,8 +49,9 @@ class SequenceBeamSearchV2(v1.SequenceBeamSearch):
     # Account for corner case where there are no finished sequences for a
     # particular batch item. In that case, return alive sequences for that batch
     # item.
-    finished_seq = tf.where(seq_cond, finished_seq, alive_seq)
-    finished_scores = tf.where(score_cond, finished_scores, alive_log_probs)
+    finished_seq = tf.compat.v2.where(seq_cond, finished_seq, alive_seq)
+    finished_scores = tf.compat.v2.where(
+        score_cond, finished_scores, alive_log_probs)
     return finished_seq, finished_scores
 
 

official/utils/logs/logger_test.py

@@ -262,14 +262,12 @@ class BenchmarkFileLoggerTest(tf.test.TestCase):
                      {"name": "batch_size", "long_value": 32})
     self.assertEqual(run_info["run_parameters"][1],
                      {"name": "dtype", "string_value": "fp16"})
-    if keras_utils.is_v2_0():
-      self.assertEqual(run_info["run_parameters"][2],
-                       {"name": "random_tensor", "string_value":
-                            "tf.Tensor(2.0, shape=(), dtype=float32)"})
-    else:
-      self.assertEqual(run_info["run_parameters"][2],
-                       {"name": "random_tensor", "string_value":
-                            "Tensor(\"Const:0\", shape=(), dtype=float32)"})
+    v1_tensor = {"name": "random_tensor", "string_value":
+                     "Tensor(\"Const:0\", shape=(), dtype=float32)"}
+    v2_tensor = {"name": "random_tensor", "string_value":
+                     "tf.Tensor(2.0, shape=(), dtype=float32)"}
+    self.assertIn(run_info["run_parameters"][2], [v1_tensor, v2_tensor])
+
 
     self.assertEqual(run_info["run_parameters"][3],
                      {"name": "resnet_size", "long_value": 50})