tf-models-nightly 2.17.0.dev20240415__py2.py3-none-any.whl → 2.17.0.dev20240417__py2.py3-none-any.whl

official/modeling/multitask/base_trainer.py

@@ -164,7 +164,7 @@ class MultiTaskBaseTrainer(orbit.StandardTrainer):
           task_metrics=self.training_metrics)
       for key, loss in losses.items():
         self.training_losses[key].update_state(loss)
+      self.global_step.assign_add(1)
 
     self.strategy.run(
         step_fn, args=(tf.nest.map_structure(next, iterator_map),))
-    self.global_step.assign_add(1)

official/modeling/multitask/interleaving_trainer.py

@@ -81,6 +81,15 @@ class MultiTaskInterleavingTrainer(base_trainer.MultiTaskBaseTrainer):
   def task_step_counter(self, name):
     return self._task_step_counters[name]
 
+  def _task_train_step(self, name):
+    """Runs one training step and updates counters."""
+    def _step_fn(inputs):
+      self._task_train_step_map[name](inputs)
+      self.global_step.assign_add(1)
+      self.task_step_counter(name).assign_add(1)
+
+    return _step_fn
+
   def train_step(self, iterator_map):
     # Sample one task to train according to a multinomial distribution
     rn = tf.random.stateless_uniform(shape=[], seed=(0, self.global_step))
@@ -96,9 +105,7 @@ class MultiTaskInterleavingTrainer(base_trainer.MultiTaskBaseTrainer):
       end = cumulative_sample_distribution[idx + 1]
       if rn >= begin and rn < end:
         self._strategy.run(
-            self._task_train_step_map[name], args=(next(iterator_map[name]),))
-        self.global_step.assign_add(1)
-        self.task_step_counter(name).assign_add(1)
+            self._task_train_step(name), args=(next(iterator_map[name]),))
 
   def train_loop_end(self):
     """Record loss and metric values per task."""

official/recommendation/uplift/metrics/poisson_metrics.py

@@ -22,6 +22,7 @@ import tensorflow as tf, tf_keras
 
 from official.recommendation.uplift import types
 from official.recommendation.uplift.metrics import loss_metric
+from official.recommendation.uplift.metrics import treatment_sliced_metric
 
 
 @tf_keras.utils.register_keras_serializable(package="Uplift")
@@ -121,3 +122,79 @@ class LogLoss(loss_metric.LossMetric):
   @classmethod
   def from_config(cls, config: dict[str, Any]) -> LogLoss:
     return cls(**config)
+
+
+def _safe_x_minus_xlogx(x: tf.Tensor) -> tf.Tensor:
+  """Computes x - x * log(x) with 0 as its continuity point when x equals 0."""
+  values = x * (1.0 - tf.math.log(x))
+  return tf.where(tf.equal(x, 0.0), tf.zeros_like(x), values)
+
+
+class LogLossMeanBaseline(tf_keras.metrics.Metric):
+  """Computes the (weighted) poisson log loss for a mean predictor."""
+
+  def __init__(
+      self,
+      compute_full_loss: bool = False,
+      slice_by_treatment: bool = True,
+      name: str = "poisson_log_loss_mean_baseline",
+      dtype: tf.DType = tf.float32,
+  ):
+    """Initializes the instance.
+
+    Args:
+      compute_full_loss: Specifies whether to compute the full poisson log loss
+        for the mean predictor or not. Defaults to `False`.
+      slice_by_treatment: Specifies whether the loss should be sliced by the
+        treatment indicator tensor. If `True`, the metric's result will return
+        the loss values sliced by the treatment group. Note that this can only
+        be set to `True` when `y_pred` is of type `TwoTowerTrainingOutputs`.
+      name: Optional name for the instance.
+      dtype: Optional data type for the instance.
+    """
+    super().__init__(name=name, dtype=dtype)
+
+    if compute_full_loss:
+      raise NotImplementedError("Full loss computation is not yet supported.")
+
+    self._compute_full_loss = compute_full_loss
+    self._slice_by_treatment = slice_by_treatment
+
+    if slice_by_treatment:
+      self._mean_label = treatment_sliced_metric.TreatmentSlicedMetric(
+          metric=tf_keras.metrics.Mean(name=name, dtype=dtype)
+      )
+    else:
+      self._mean_label = tf_keras.metrics.Mean(name=name, dtype=dtype)
+
+  def update_state(
+      self,
+      y_true: tf.Tensor,
+      y_pred: types.TwoTowerTrainingOutputs | tf.Tensor | None = None,
+      sample_weight: tf.Tensor | None = None,
+  ):
+    is_treatment = {}
+    if self._slice_by_treatment:
+      if not isinstance(y_pred, types.TwoTowerTrainingOutputs):
+        raise ValueError(
+            "`slice_by_treatment` must be set to `False` when `y_pred` is not"
+            " of type `TwoTowerTrainingOutputs`."
+        )
+      is_treatment["is_treatment"] = y_pred.is_treatment
+
+    self._mean_label.update_state(
+        y_true, sample_weight=sample_weight, **is_treatment
+    )
+
+  def result(self) -> tf.Tensor | dict[str, tf.Tensor]:
+    return tf.nest.map_structure(_safe_x_minus_xlogx, self._mean_label.result())
+
+  def get_config(self) -> dict[str, Any]:
+    config = super().get_config()
+    config["compute_full_loss"] = self._compute_full_loss
+    config["slice_by_treatment"] = self._slice_by_treatment
+    return config
+
+  @classmethod
+  def from_config(cls, config: dict[str, Any]) -> LogLossMeanBaseline:
+    return cls(**config)

official/recommendation/uplift/metrics/poisson_metrics_test.py

@@ -23,7 +23,8 @@ from official.recommendation.uplift.metrics import poisson_metrics
 
 
 def _get_two_tower_outputs(
-    true_logits: tf.Tensor, is_treatment: tf.Tensor
+    is_treatment: tf.Tensor,
+    true_logits: tf.Tensor | None = None,
 ) -> types.TwoTowerTrainingOutputs:
   # Only the true_logits and is_treatment tensors are needed for testing.
   return types.TwoTowerTrainingOutputs(
@@ -33,7 +34,9 @@ def _get_two_tower_outputs(
       uplift=tf.ones_like(is_treatment),
       control_logits=tf.ones_like(is_treatment),
       treatment_logits=tf.ones_like(is_treatment),
-      true_logits=true_logits,
+      true_logits=(
+          true_logits if true_logits is not None else tf.ones_like(is_treatment)
+      ),
       true_predictions=tf.ones_like(is_treatment),
       is_treatment=is_treatment,
   )
@@ -177,5 +180,81 @@ class LogLossTest(keras_test_case.KerasTestCase, parameterized.TestCase):
     )
 
 
+class LogLossMeanBaselineTest(
+    keras_test_case.KerasTestCase, parameterized.TestCase
+):
+
+  @parameterized.named_parameters(
+      {
+          "testcase_name": "label_zero",
+          "expected_loss": 0.0,
+          "y_true": tf.constant([0], dtype=tf.float32),
+      },
+      {
+          "testcase_name": "small_positive_label",
+          "expected_loss": 0.0,
+          "y_true": tf.constant([1e-10], dtype=tf.float32),
+      },
+      {
+          "testcase_name": "label_one",
+          "expected_loss": 1.0,
+          "y_true": tf.constant([1], dtype=tf.float32),
+      },
+      {
+          "testcase_name": "weighted_loss",
+          "expected_loss": 1.0,
+          "y_true": tf.constant([[0], [1]], dtype=tf.float32),
+          "sample_weight": tf.constant([[0], [1]], dtype=tf.float32),
+      },
+      {
+          "testcase_name": "two_tower_outputs",
+          "expected_loss": 0.5 - 0.5 * tf.math.log(0.5),
+          "y_true": tf.constant([[0], [1]], dtype=tf.float32),
+          "y_pred": _get_two_tower_outputs(
+              is_treatment=tf.constant([[0], [1]], dtype=tf.float32),
+          ),
+      },
+      {
+          "testcase_name": "two_tower_outputs_sliced_loss",
+          "expected_loss": {
+              "loss": 0.5 - 0.5 * tf.math.log(0.5),
+              "loss/control": 0.0,
+              "loss/treatment": 1.0,
+          },
+          "y_true": tf.constant([[0], [1]], dtype=tf.float32),
+          "y_pred": _get_two_tower_outputs(
+              is_treatment=tf.constant([[0], [1]], dtype=tf.float32),
+          ),
+          "slice_by_treatment": True,
+      },
+  )
+  def test_metric_computes_correct_loss(
+      self,
+      expected_loss: tf.Tensor,
+      y_true: tf.Tensor,
+      y_pred: types.TwoTowerTrainingOutputs | tf.Tensor | None = None,
+      sample_weight: tf.Tensor | None = None,
+      slice_by_treatment: bool = False,
+  ):
+    metric = poisson_metrics.LogLossMeanBaseline(
+        slice_by_treatment=slice_by_treatment, name="loss"
+    )
+    metric.update_state(y_true, y_pred, sample_weight=sample_weight)
+    self.assertAllClose(expected_loss, metric.result())
+
+  def test_negative_label_returns_nan_loss(self):
+    metric = poisson_metrics.LogLossMeanBaseline(slice_by_treatment=False)
+    metric.update_state(tf.constant([-1.0]))
+    self.assertTrue(tf.math.is_nan(metric.result()).numpy().item())
+
+  def test_metric_is_configurable(self):
+    metric = poisson_metrics.LogLossMeanBaseline(slice_by_treatment=False)
+    self.assertLayerConfigurable(
+        layer=metric,
+        y_true=tf.constant([[0], [0], [2], [7]], dtype=tf.float32),
+        serializable=True,
+    )
+
+
 if __name__ == "__main__":
   tf.test.main()

official/vision/modeling/backbones/__init__.py

@@ -14,6 +14,7 @@
 
 """Backbones package definition."""
 
+from official.projects.maskconver.modeling.resnet_unet import ResNetUNet
 from official.vision.modeling.backbones.efficientnet import EfficientNet
 from official.vision.modeling.backbones.mobiledet import MobileDet
 from official.vision.modeling.backbones.mobilenet import MobileNet

{tf_models_nightly-2.17.0.dev20240415.dist-info → tf_models_nightly-2.17.0.dev20240417.dist-info}/METADATA

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

{tf_models_nightly-2.17.0.dev20240415.dist-info → tf_models_nightly-2.17.0.dev20240417.dist-info}/RECORD

@@ -213,12 +213,12 @@ official/modeling/hyperparams/params_dict.py,sha256=63fftQdUlycgJErxcyIj7655zL57
 official/modeling/hyperparams/params_dict_test.py,sha256=WPX-VU7L3JVjS42b7BWe77QxP1kdwch05Ib7LvNOTYs,14673
 official/modeling/multitask/__init__.py,sha256=7oiypy0N82PDw9aSdcJBLVoGTd_oRSUOdvuJhMv4leQ,609
 official/modeling/multitask/base_model.py,sha256=QI8qb8ipj75IUj6bKNjcAFHPjeqmNjqHr7nUbPd6a-o,1946
-official/modeling/multitask/base_trainer.py,sha256=niGf1XLtsYyvJE7NFYXphTmIPwKlJbpS34PGK3VgQFo,5856
+official/modeling/multitask/base_trainer.py,sha256=83cLDajiyS2lJPMhllTdIsKXqiVTFLDaGZaherTPCa8,5858
 official/modeling/multitask/base_trainer_test.py,sha256=qJ7z4kid2XAX6hOIvUHa7dwqxouemMekS9ZXhPjWW9w,3663
 official/modeling/multitask/configs.py,sha256=ZO2waQrMn9CAgyFpsmeQvplCF5VeXz7tCPmIuy5jvlc,3164
 official/modeling/multitask/evaluator.py,sha256=spDm2X8EX62qsxI2ehVjrkIKoo-omQQOYcAVKZNgxHc,6078
 official/modeling/multitask/evaluator_test.py,sha256=vU-q-gM7GqiMqE5zbBnOT8mPFhQmHjniMyNnwganhso,4643
-official/modeling/multitask/interleaving_trainer.py,sha256=ZZHKsqbJKLqvwtgy-PUv_S_8bDG0MhJDNwWICY_IF6Q,4458
+official/modeling/multitask/interleaving_trainer.py,sha256=f111ZhknyS34hpP0FfdWjX3_iiLViHfBd0VSuC715s0,4635
 official/modeling/multitask/interleaving_trainer_test.py,sha256=MeQQxpcinPTQuTrAcITjwHa2bAj-XCBCqYsrbxPBus8,4305
 official/modeling/multitask/multitask.py,sha256=DV-ysfhPiIZgsrzZNylsPBxKNBf_xzPxJYjF4buWVgE,5948
 official/modeling/multitask/task_sampler.py,sha256=SGVVdjMb5oG4vnCczpfdgBtbsdsXiyBLl9si_0V6nko,4897
@@ -915,8 +915,8 @@ official/recommendation/uplift/metrics/label_variance_test.py,sha256=k0mdEU1WU53
 official/recommendation/uplift/metrics/loss_metric.py,sha256=gYZdnTsuL_2q1FZuPip-DaWxt_Q-02YYaePyMBVNx7w,7344
 official/recommendation/uplift/metrics/loss_metric_test.py,sha256=48rQG8bKFdy0xBFjoOLXKRUlYpCEyAzSmPOFoF7FX94,16021
 official/recommendation/uplift/metrics/metric_configs.py,sha256=Z-r79orE4EycQ5TJ7xdI5LhjOHT3wzChYyDxcxGqLXk,1670
-official/recommendation/uplift/metrics/poisson_metrics.py,sha256=LJnovpST0H9kFGu-ziDstWOVlAYARLo9oPLDTjzrdu4,4623
-official/recommendation/uplift/metrics/poisson_metrics_test.py,sha256=Kd8CuQeEBlxRklA-7mGKHcUD0CyskE1S3cJqk6mEvv4,6756
+official/recommendation/uplift/metrics/poisson_metrics.py,sha256=1a0sHwdOLUVhmsdpnZIQxAWkt3QAPWS9TZCK6P9H7ug,7443
+official/recommendation/uplift/metrics/poisson_metrics_test.py,sha256=sLx4RjtvJiD8kRkXBDfsUANK0P5LSMfTTPj9GF73wOw,9473
 official/recommendation/uplift/metrics/sliced_metric.py,sha256=uhvzudOWtMNKZ0avwGhX-37UELR9Cq9b4C0g8erBkXw,8688
 official/recommendation/uplift/metrics/sliced_metric_test.py,sha256=bhVGyI1tOkFkVOtruJo3p6XopDFyG1JW5qdZm9-RqeU,12248
 official/recommendation/uplift/metrics/treatment_fraction.py,sha256=WHrKfsN42xU7S-pK99xEVpVtd3zLD7UidLT1K8vgIn4,2757
@@ -1042,7 +1042,7 @@ official/vision/modeling/segmentation_model.py,sha256=BrH3w0yZ60uwyEGP08iYt0rAWV
 official/vision/modeling/segmentation_model_test.py,sha256=fimmpt0pFMcX214Vuf0C4NcRdWLLB9F2IJeiuTFUO54,2817
 official/vision/modeling/video_classification_model.py,sha256=RKrT22D5yW435S9AyYlgOFCIi0knbmZFvgwh_5JubPI,4713
 official/vision/modeling/video_classification_model_test.py,sha256=HL0hVwOFJCxSjDNFVK42MwzuoDx5r4rvraEbBXVkX8c,3271
-official/vision/modeling/backbones/__init__.py,sha256=YqnYYFcfa0LAM6vB8LQvBIuTRw7RSpWw5zIm5vrve3E,1329
+official/vision/modeling/backbones/__init__.py,sha256=21iuq1HPa3KcN-6ljjRlUyklRV8Ynyc3Cz47XivHWMw,1402
 official/vision/modeling/backbones/efficientnet.py,sha256=j716OGSpkzgpDA4jV-Hk73mAGPULrPwlwHuS_9j40bE,12438
 official/vision/modeling/backbones/efficientnet_test.py,sha256=TYsUieiLrEU5913s3Yxhv-9eaolQK_kfzJm77lyCF0M,3762
 official/vision/modeling/backbones/factory.py,sha256=coJKJpPMhgM9gAc2Q7I5_CuzAaHZNJwPcvGbaUYp8gU,3504
@@ -1206,9 +1206,9 @@ tensorflow_models/tensorflow_models_test.py,sha256=nc6A9K53OGqF25xN5St8EiWvdVbda
 tensorflow_models/nlp/__init__.py,sha256=4tA5Pf4qaFwT-fIFOpX7x7FHJpnyJT-5UgOeFYTyMlc,807
 tensorflow_models/uplift/__init__.py,sha256=mqfa55gweOdpKoaQyid4A_4u7xw__FcQeSIF0k_pYmI,999
 tensorflow_models/vision/__init__.py,sha256=zBorY_v5xva1uI-qxhZO3Qh-Dii-Suq6wEYh6hKHDfc,833
-tf_models_nightly-2.17.0.dev20240415.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.17.0.dev20240415.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.17.0.dev20240415.dist-info/METADATA,sha256=5gRdeXttXibMvRXrmtEuZtlJJBvHIgrltUkv5myi9T8,1432
-tf_models_nightly-2.17.0.dev20240415.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.17.0.dev20240415.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.17.0.dev20240415.dist-info/RECORD,,
+tf_models_nightly-2.17.0.dev20240417.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.17.0.dev20240417.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.17.0.dev20240417.dist-info/METADATA,sha256=Sft4CH1GT2IMtgKOIyKA8O9FFE9nlzhWqHn2oHYV-G4,1432
+tf_models_nightly-2.17.0.dev20240417.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.17.0.dev20240417.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.17.0.dev20240417.dist-info/RECORD,,