Correct imports of keras loss utils

PiperOrigin-RevId: 486795765

Author: tensorflower-gardener

tensorflow_privacy/privacy/estimators/binary_class_head.py

@@ -14,7 +14,6 @@
 """Binary class head for Estimator that allow integration with TF Privacy."""
 
 import tensorflow as tf
-from tensorflow.python.keras.utils import losses_utils  # pylint: disable=g-direct-tensorflow-import
 from tensorflow_estimator.python.estimator import model_fn
 from tensorflow_estimator.python.estimator.canned import prediction_keys
 from tensorflow_estimator.python.estimator.export import export_output
@@ -55,7 +54,7 @@ def loss(self,
       labels = self._processed_labels(logits, labels)
       unweighted_loss, weights = self._unweighted_loss_and_weights(
           logits, labels, features)
-      vector_training_loss = losses_utils.compute_weighted_loss(
+      vector_training_loss = tf.keras.__internal__.losses.compute_weighted_loss(
           unweighted_loss,
           sample_weight=weights,
           reduction=tf.keras.losses.Reduction.NONE)

tensorflow_privacy/privacy/estimators/multi_class_head.py

@@ -14,7 +14,6 @@
 """Multiclass head for Estimator that allow integration with TF Privacy."""
 
 import tensorflow as tf
-from tensorflow.python.keras.utils import losses_utils  # pylint: disable=g-direct-tensorflow-import
 from tensorflow_estimator.python.estimator import model_fn
 from tensorflow_estimator.python.estimator.canned import prediction_keys
 from tensorflow_estimator.python.estimator.export import export_output
@@ -30,14 +29,14 @@ def __init__(self,
                n_classes,
                weight_column=None,
                label_vocabulary=None,
-               loss_reduction=losses_utils.ReductionV2.SUM_OVER_BATCH_SIZE,
+               loss_reduction=tf.keras.losses.Reduction.SUM_OVER_BATCH_SIZE,
                loss_fn=None,
                name=None):
     super().__init__(
         n_classes=n_classes,
         weight_column=weight_column,
         label_vocabulary=label_vocabulary,
-        loss_reduction=losses_utils.ReductionV2.SUM_OVER_BATCH_SIZE,
+        loss_reduction=tf.keras.losses.Reduction.SUM_OVER_BATCH_SIZE,
         loss_fn=loss_fn,
         name=name)
 
@@ -55,7 +54,7 @@ def loss(self,
       labels = self._processed_labels(logits, labels)
       unweighted_loss, weights = self._unweighted_loss_and_weights(
           logits, labels, features)
-      vector_training_loss = losses_utils.compute_weighted_loss(
+      vector_training_loss = tf.keras.__internal__.losses.compute_weighted_loss(
           unweighted_loss,
           sample_weight=weights,
           reduction=tf.keras.losses.Reduction.NONE)

tensorflow_privacy/privacy/estimators/multi_label_head.py

@@ -14,7 +14,6 @@
 """Multiclass head for Estimator that allow integration with TF Privacy."""
 
 import tensorflow as tf
-from tensorflow.python.keras.utils import losses_utils  # pylint: disable=g-direct-tensorflow-import
 from tensorflow_estimator.python.estimator import model_fn
 from tensorflow_estimator.python.estimator.canned import prediction_keys
 from tensorflow_estimator.python.estimator.export import export_output
@@ -61,7 +60,7 @@ def loss(self,
       labels = self._processed_labels(logits, labels)
       unweighted_loss, weights = self._unweighted_loss_and_weights(
           logits, labels, features)
-      vector_training_loss = losses_utils.compute_weighted_loss(
+      vector_training_loss = tf.keras.__internal__.losses.compute_weighted_loss(
           unweighted_loss,
           sample_weight=weights,
           reduction=tf.keras.losses.Reduction.NONE)