Avoid importing private ObjectIdentitySet class

PiperOrigin-RevId: 266848625

official/bert/model_training_utils.py

@@ -23,7 +23,6 @@ import os
 
 from absl import logging
 import tensorflow as tf
-from tensorflow.python.util import object_identity
 from official.utils.misc import distribution_utils
 from official.utils.misc import tpu_lib
 
@@ -243,8 +242,7 @@ def run_customized_training_loop(
             scaled_loss = optimizer.get_scaled_loss(loss)
 
         # De-dupes variables due to keras tracking issues.
-        tvars = list(
-            object_identity.ObjectIdentitySet(model.trainable_variables))
+        tvars = list({id(v): v for v in model.trainable_variables}.values())
         if use_float16:
           scaled_grads = tape.gradient(scaled_loss, tvars)
           grads = optimizer.get_unscaled_gradients(scaled_grads)

official/transformer/v2/transformer_main.py

@@ -30,8 +30,6 @@ from absl import flags
 from absl import logging
 import tensorflow as tf
 
-from tensorflow.python.util import object_identity
-
 # pylint: disable=g-bad-import-order
 from official.transformer import compute_bleu
 from official.transformer.utils import tokenizer
@@ -271,8 +269,7 @@ class TransformerTask(object):
           scaled_loss = loss / self.distribution_strategy.num_replicas_in_sync
 
         # De-dupes variables due to keras tracking issues.
-        tvars = list(
-            object_identity.ObjectIdentitySet(model.trainable_variables))
+        tvars = list({id(v): v for v in model.trainable_variables}.values())
         grads = tape.gradient(scaled_loss, tvars)
         opt.apply_gradients(zip(grads, tvars))
         # For reporting, the metric takes the mean of losses.