tf-models-nightly 2.17.0.dev20240409__py2.py3-none-any.whl → 2.17.0.dev20240410__py2.py3-none-any.whl

official/recommendation/uplift/metrics/loss_metric.py

@@ -62,6 +62,7 @@ class LossMetric(tf_keras.metrics.Metric):
           Callable[[tf.Tensor, tf.Tensor], tf.Tensor] | tf_keras.metrics.Metric
       ),
       from_logits: bool = True,
+      slice_by_treatment: bool = True,
       name: str = "loss",
       dtype: tf.DType = tf.float32,
       **loss_fn_kwargs,
@@ -76,6 +77,10 @@ class LossMetric(tf_keras.metrics.Metric):
       from_logits: Specifies whether the true logits or true predictions should
         be used from the model outputs to compute the loss. Defaults to using
         the true logits.
+      slice_by_treatment: Specifies whether the loss should be sliced by the
+        treatment indicator tensor. If `True`, `loss_fn` will be wrapped in a
+        `TreatmentSlicedMetric` to report the loss values sliced by the
+        treatment group.
       name: Optional name for the instance. If `loss_fn` is a Keras metric then
         its name will be used instead.
       dtype: Optional data type for the instance. If `loss_fn` is a Keras metric
@@ -99,6 +104,7 @@ class LossMetric(tf_keras.metrics.Metric):
     self._loss_fn = loss_fn
     self._from_logits = from_logits
     self._loss_fn_kwargs = loss_fn_kwargs
+    self._slice_by_treatment = slice_by_treatment
 
     if isinstance(loss_fn, tf_keras.metrics.Metric):
       metric_from_logits = loss_fn.get_config().get("from_logits", from_logits)
@@ -108,20 +114,17 @@ class LossMetric(tf_keras.metrics.Metric):
             " the `from_logits` value passed to the `loss_fn` metric"
             f" ({metric_from_logits}). Ensure that they have the same value."
         )
-
-      self._treatment_sliced_loss = (
-          treatment_sliced_metric.TreatmentSlicedMetric(loss_fn)
-      )
+      loss_metric = loss_fn
 
     else:
       if "from_logits" in inspect.signature(loss_fn).parameters:
         self._loss_fn_kwargs.update({"from_logits": from_logits})
+      loss_metric = tf_keras.metrics.Mean(name=name, dtype=dtype)
 
-      self._treatment_sliced_loss = (
-          treatment_sliced_metric.TreatmentSlicedMetric(
-              tf_keras.metrics.Mean(name=name, dtype=dtype)
-          )
-      )
+    if slice_by_treatment:
+      self._loss = treatment_sliced_metric.TreatmentSlicedMetric(loss_metric)
+    else:
+      self._loss = loss_metric
 
   def update_state(
       self,
@@ -151,27 +154,32 @@ class LossMetric(tf_keras.metrics.Metric):
 
     pred = y_pred.true_logits if self._from_logits else y_pred.true_predictions
 
+    is_treatment = {}
+    if self._slice_by_treatment:
+      is_treatment["is_treatment"] = y_pred.is_treatment
+
     if isinstance(self._loss_fn, tf_keras.metrics.Metric):
-      self._treatment_sliced_loss.update_state(
+      self._loss.update_state(
           y_true,
           y_pred=pred,
-          is_treatment=y_pred.is_treatment,
           sample_weight=sample_weight,
+          **is_treatment,
       )
     else:
-      self._treatment_sliced_loss.update_state(
+      self._loss.update_state(
           values=self._loss_fn(y_true, pred, **self._loss_fn_kwargs),
-          is_treatment=y_pred.is_treatment,
           sample_weight=sample_weight,
+          **is_treatment,
       )
 
-  def result(self) -> dict[str, tf.Tensor]:
-    return self._treatment_sliced_loss.result()
+  def result(self) -> tf.Tensor | dict[str, tf.Tensor]:
+    return self._loss.result()
 
   def get_config(self) -> dict[str, Any]:
     config = super().get_config()
     config["loss_fn"] = tf_keras.utils.serialize_keras_object(self._loss_fn)
     config["from_logits"] = self._from_logits
+    config["slice_by_treatment"] = self._slice_by_treatment
     config.update(self._loss_fn_kwargs)
     return config
 
@@ -180,4 +188,4 @@ class LossMetric(tf_keras.metrics.Metric):
     config["loss_fn"] = tf_keras.utils.deserialize_keras_object(
         config["loss_fn"]
     )
-    return LossMetric(**config)
+    return cls(**config)

official/recommendation/uplift/metrics/loss_metric_test.py

@@ -25,9 +25,7 @@ from official.recommendation.uplift import types
 from official.recommendation.uplift.metrics import loss_metric
 
 
-class LossMetricTest(
-    keras_test_case.KerasTestCase, parameterized.TestCase
-):
+class LossMetricTest(keras_test_case.KerasTestCase, parameterized.TestCase):
 
   def _get_outputs(
       self,
@@ -236,6 +234,28 @@ class LossMetricTest(
               "loss/treatment": 0.0,
           },
       },
+      {
+          "testcase_name": "no_treatment_slice",
+          "loss_fn": tf_keras.losses.binary_crossentropy,
+          "from_logits": True,
+          "y_true": tf.constant([[0.0, 1.0]]),
+          "y_pred": tf.constant([[0.0, 1.0]]),
+          "is_treatment": tf.constant([[0], [0]]),
+          "sample_weight": None,
+          "expected_losses": 0.50320446,
+          "slice_by_treatment": False,
+      },
+      {
+          "testcase_name": "no_treatment_slice_metric",
+          "loss_fn": tf_keras.metrics.BinaryCrossentropy(from_logits=False),
+          "from_logits": False,
+          "y_true": tf.constant([[0.0, 1.0]]),
+          "y_pred": tf.constant([[0.0, 1.0]]),
+          "is_treatment": tf.constant([[0], [0]]),
+          "sample_weight": None,
+          "expected_losses": 0,
+          "slice_by_treatment": False,
+      },
   )
   def test_metric_computes_sliced_losses(
       self,
@@ -245,7 +265,8 @@ class LossMetricTest(
       y_pred: tf.Tensor,
       is_treatment: tf.Tensor,
       sample_weight: tf.Tensor | None,
-      expected_losses: dict[str, float],
+      expected_losses: float | dict[str, float],
+      slice_by_treatment: bool = True,
   ):
     if from_logits:
       true_logits = y_pred
@@ -254,7 +275,11 @@ class LossMetricTest(
       true_logits = tf.zeros_like(y_pred)  # Irrelevant for testing.
       true_predictions = y_pred
 
-    metric = loss_metric.LossMetric(loss_fn=loss_fn, from_logits=from_logits)
+    metric = loss_metric.LossMetric(
+        loss_fn=loss_fn,
+        from_logits=from_logits,
+        slice_by_treatment=slice_by_treatment,
+    )
     outputs = self._get_outputs(
         true_logits=true_logits,
         true_predictions=true_predictions,
@@ -335,17 +360,28 @@ class LossMetricTest(
     metric.reset_states()
     self.assertEqual(expected_initial_result, metric.result())
 
-  @parameterized.parameters(
-      tf_keras.losses.binary_crossentropy,
-      tf_keras.metrics.BinaryCrossentropy(from_logits=True, name="bce_loss"),
+  @parameterized.product(
+      loss_fn=(
+          tf_keras.losses.binary_crossentropy,
+          tf_keras.metrics.BinaryCrossentropy(
+              from_logits=True, name="bce_loss"
+          ),
+      ),
+      slice_by_treatment=(True, False),
   )
   def test_metric_is_configurable(
       self,
       loss_fn: (
           Callable[[tf.Tensor, tf.Tensor], tf.Tensor] | tf_keras.metrics.Metric
       ),
+      slice_by_treatment: bool,
   ):
-    metric = loss_metric.LossMetric(loss_fn, from_logits=True, name="bce_loss")
+    metric = loss_metric.LossMetric(
+        loss_fn,
+        from_logits=True,
+        slice_by_treatment=slice_by_treatment,
+        name="bce_loss",
+    )
     self.assertLayerConfigurable(
         layer=metric,
         y_true=tf.constant([[1], [1], [0]]),

{tf_models_nightly-2.17.0.dev20240409.dist-info → tf_models_nightly-2.17.0.dev20240410.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tf-models-nightly
-Version: 2.17.0.dev20240409
+Version: 2.17.0.dev20240410
 Summary: TensorFlow Official Models
 Home-page: https://github.com/tensorflow/models
 Author: Google Inc.

{tf_models_nightly-2.17.0.dev20240409.dist-info → tf_models_nightly-2.17.0.dev20240410.dist-info}/RECORD

@@ -912,8 +912,8 @@ official/recommendation/uplift/metrics/label_mean.py,sha256=ECaes7FZmsksnwySn7jf
 official/recommendation/uplift/metrics/label_mean_test.py,sha256=b_d3lNlpkDm2xKLUkxfiXeQg7pjL8HNx7y9NaYarpV0,7083
 official/recommendation/uplift/metrics/label_variance.py,sha256=9DCl42BJkehxfWD3pSbZnRNvwfhVM6VyHwivGdaU72s,3610
 official/recommendation/uplift/metrics/label_variance_test.py,sha256=k0mdEU1WU53-HIEO5HGtfp1MleifD-h4bZNKtTvM3Ws,7681
-official/recommendation/uplift/metrics/loss_metric.py,sha256=8Lfi4FNR6uj8eLSdT5t0XjEzTZtwL-xv5e5KNGzrYWU,6498
-official/recommendation/uplift/metrics/loss_metric_test.py,sha256=yropE1t8PblsJ_kJEno5ratkY_ka81PHgcfroWtRKVI,13768
+official/recommendation/uplift/metrics/loss_metric.py,sha256=qWMcLk1JmLkru5T9qcqO2UbB_AJQKjQ8uUSAxjR1l58,6883
+official/recommendation/uplift/metrics/loss_metric_test.py,sha256=QIUujZvUud-A2itI2vs2KpzWfLwbq2XDd9H9CFq5D4Y,14929
 official/recommendation/uplift/metrics/metric_configs.py,sha256=Z-r79orE4EycQ5TJ7xdI5LhjOHT3wzChYyDxcxGqLXk,1670
 official/recommendation/uplift/metrics/sliced_metric.py,sha256=O2I2apZK6IfOQK9Q_mgSiTUCnGokczp4e14zrrYNeRU,8564
 official/recommendation/uplift/metrics/sliced_metric_test.py,sha256=dhY41X8lqT_WW04XLjyjDerZwujEBGeTXtxf4NkYThw,11359
@@ -1203,9 +1203,9 @@ tensorflow_models/__init__.py,sha256=etxw45SHxuwFCRX5qGxGMP83II0JfJulzNl5GSNJvhw
 tensorflow_models/tensorflow_models_test.py,sha256=AxUYUdiQn416UR7jg0h6rmv688esvlKDfpyDCIQkF18,1395
 tensorflow_models/nlp/__init__.py,sha256=4tA5Pf4qaFwT-fIFOpX7x7FHJpnyJT-5UgOeFYTyMlc,807
 tensorflow_models/vision/__init__.py,sha256=zBorY_v5xva1uI-qxhZO3Qh-Dii-Suq6wEYh6hKHDfc,833
-tf_models_nightly-2.17.0.dev20240409.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.17.0.dev20240409.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.17.0.dev20240409.dist-info/METADATA,sha256=738N3KUX3sKTeF-7mDL3_2r1htitc4BCT9A4QMgB9U4,1432
-tf_models_nightly-2.17.0.dev20240409.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.17.0.dev20240409.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.17.0.dev20240409.dist-info/RECORD,,
+tf_models_nightly-2.17.0.dev20240410.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.17.0.dev20240410.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.17.0.dev20240410.dist-info/METADATA,sha256=Bh56e1__Pn_yvpm9QMq4h12nKucvnKAftuk0ElQqn_w,1432
+tf_models_nightly-2.17.0.dev20240410.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.17.0.dev20240410.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.17.0.dev20240410.dist-info/RECORD,,