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

official/modeling/multitask/interleaving_trainer.py

@@ -43,6 +43,12 @@ class MultiTaskInterleavingTrainer(base_trainer.MultiTaskBaseTrainer):
         trainer_options=trainer_options)
     self._task_sampler = task_sampler
 
+    # TODO(haozhangthu): Add taskwise step counter to train_loop_end for logging
+    # on TensorBoard.
+    self._task_step_counters = {
+        name: orbit.utils.create_global_step() for name in self.multi_task.tasks
+    }
+
     # Build per task train step.
     def _get_task_step(task_name, task):
 
@@ -57,6 +63,8 @@ class MultiTaskInterleavingTrainer(base_trainer.MultiTaskBaseTrainer):
             optimizer=self.optimizer,
             metrics=self.training_metrics[task_name])
         self.training_losses[task_name].update_state(task_logs[task.loss])
+        self.global_step.assign_add(1)
+        self.task_step_counter(task_name).assign_add(1)
 
       return step_fn
 
@@ -65,12 +73,6 @@ class MultiTaskInterleavingTrainer(base_trainer.MultiTaskBaseTrainer):
         for name, task in self.multi_task.tasks.items()
     }
 
-    # TODO(haozhangthu): Add taskwise step counter to train_loop_end for logging
-    # on TensorBoard.
-    self._task_step_counters = {
-        name: orbit.utils.create_global_step() for name in self.multi_task.tasks
-    }
-
     # If the new Keras optimizer is used, we require all model variables are
     # created before the training and let the optimizer to create the slot
     # variable all together.
@@ -97,8 +99,6 @@ class MultiTaskInterleavingTrainer(base_trainer.MultiTaskBaseTrainer):
       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)
 
   def train_loop_end(self):
     """Record loss and metric values per task."""

official/recommendation/uplift/metrics/loss_metric.py

@@ -19,6 +19,7 @@ from __future__ import annotations
 import inspect
 from typing import Any, Callable
 
+import numpy as np
 import tensorflow as tf, tf_keras
 
 from official.recommendation.uplift import types
@@ -29,9 +30,6 @@ from official.recommendation.uplift.metrics import treatment_sliced_metric
 class LossMetric(tf_keras.metrics.Metric):
   """Computes a loss sliced by treatment group.
 
-  Note that the prediction tensor is expected to be of type
-  `TwoTowerTrainingOutputs`.
-
   Example standalone usage:
 
   >>> sliced_loss = LossMetric(tf_keras.losses.mean_squared_error)
@@ -74,9 +72,10 @@ class LossMetric(tf_keras.metrics.Metric):
         `__call__(y_true: tf,Tensor, y_pred: tf.Tensor, **loss_fn_kwargs)`. Note
         that the `loss_fn_kwargs` will not be passed to the `__call__` method if
         `loss_fn` is a Keras 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.
+      from_logits: When `y_pred` is of type `TwoTowerTrainingOutputs`, specifies
+        whether the true logits or true predictions should be used to compute
+        the loss (defaults to using the true logits). Othwerwise, this argument
+        will be ignored if `y_pred` is of type `tf.Tensor`.
       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
@@ -129,16 +128,16 @@ class LossMetric(tf_keras.metrics.Metric):
   def update_state(
       self,
       y_true: tf.Tensor,
-      y_pred: types.TwoTowerTrainingOutputs,
+      y_pred: types.TwoTowerTrainingOutputs | tf.Tensor | np.ndarray,
       sample_weight: tf.Tensor | None = None,
   ):
     """Updates the overall, control and treatment losses.
 
     Args:
       y_true: A `tf.Tensor` with the targets.
-      y_pred: Two tower training outputs. The treatment indicator tensor is used
-        to slice the true logits or true predictions into control and treatment
-        losses.
+      y_pred: Model outputs. If of type `TwoTowerTrainingOutputs`, the treatment
+        indicator tensor is used to slice the true logits or true predictions
+        into control and treatment losses.
       sample_weight: Optional sample weight to compute weighted losses. If
         given, the sample weight will also be sliced by the treatment indicator
         tensor to compute the weighted control and treatment losses.
@@ -146,14 +145,23 @@ class LossMetric(tf_keras.metrics.Metric):
     Raises:
       TypeError: if `y_pred` is not of type `TwoTowerTrainingOutputs`.
     """
-    if not isinstance(y_pred, types.TwoTowerTrainingOutputs):
+    if isinstance(y_pred, (tf.Tensor, np.ndarray)):
+      if self._slice_by_treatment:
+        raise ValueError(
+            "`slice_by_treatment` must be False when y_pred is a `tf.Tensor` or"
+            " `np.ndarray`."
+        )
+      pred = y_pred
+    elif isinstance(y_pred, types.TwoTowerTrainingOutputs):
+      pred = (
+          y_pred.true_logits if self._from_logits else y_pred.true_predictions
+      )
+    else:
       raise TypeError(
-          "y_pred must be of type `TwoTowerTrainingOutputs` but got type"
-          f" {type(y_pred)} instead."
+          "y_pred must be of type `TwoTowerTrainingOutputs`, `tf.Tensor` or"
+          f" `np.ndarray` but got type {type(y_pred)} instead."
       )
 
-    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

official/recommendation/uplift/metrics/loss_metric_test.py

@@ -288,6 +288,19 @@ class LossMetricTest(keras_test_case.KerasTestCase, parameterized.TestCase):
     metric(y_true, outputs, sample_weight=sample_weight)
     self.assertEqual(expected_losses, metric.result())
 
+  def test_metric_with_y_pred_tensor(self):
+    y_true = tf.constant([[0], [0], [2], [7]])
+    y_pred = tf.constant([[1], [2], [3], [4]])
+    sample_weight = tf.constant([[0.5], [0.5], [0.7], [1.8]])
+
+    metric = loss_metric.LossMetric(
+        loss_fn=tf_keras.metrics.mae, slice_by_treatment=False
+    )
+    metric(y_true, y_pred, sample_weight)
+
+    expected_loss = np.average([1, 2, 1, 3], weights=[0.5, 0.5, 0.7, 1.8])
+    self.assertAllClose(expected_loss, metric.result())
+
   def test_multiple_update_batches_returns_aggregated_sliced_losses(self):
     metric = loss_metric.LossMetric(
         loss_fn=tf_keras.losses.mean_absolute_error,
@@ -397,8 +410,26 @@ class LossMetricTest(keras_test_case.KerasTestCase, parameterized.TestCase):
     metric = loss_metric.LossMetric(
         tf_keras.metrics.mean_absolute_percentage_error
     )
+    y_true = tf.ones((3, 1))
+    y_pred = types.TwoTowerNetworkOutputs(
+        shared_embedding=tf.ones((3, 5)),
+        control_logits=tf.ones((3, 1)),
+        treatment_logits=tf.ones((3, 1)),
+    )
+    with self.assertRaisesRegex(
+        TypeError,
+        "y_pred must be of type `TwoTowerTrainingOutputs`, `tf.Tensor` or"
+        " `np.ndarray`",
+    ):
+      metric.update_state(y_true=y_true, y_pred=y_pred)
+
+  def test_slice_by_treatment_with_y_pred_tensor_raises_error(self):
+    metric = loss_metric.LossMetric(
+        tf_keras.metrics.mae, slice_by_treatment=True
+    )
     with self.assertRaisesRegex(
-        TypeError, "y_pred must be of type `TwoTowerTrainingOutputs`"
+        ValueError,
+        "`slice_by_treatment` must be False when y_pred is a `tf.Tensor`.",
     ):
       metric.update_state(y_true=tf.ones((3, 1)), y_pred=tf.ones((3, 1)))
 

tensorflow_models/__init__.py

@@ -15,6 +15,7 @@
 """TensorFlow Models Libraries."""
 # pylint: disable=wildcard-import
 from tensorflow_models import nlp
+from tensorflow_models import uplift
 from tensorflow_models import vision
 
 from official import core

tensorflow_models/tensorflow_models_test.py

@@ -35,6 +35,17 @@ class TensorflowModelsTest(tf.test.TestCase):
     _ = tfm.optimization.LinearWarmup(
         after_warmup_lr_sched=0.0, warmup_steps=10, warmup_learning_rate=0.1)
 
+  def testUpliftImports(self):
+    _ = tfm.uplift.keys.TwoTowerOutputKeys.CONTROL_PREDICTIONS
+    _ = tfm.uplift.types.TwoTowerNetworkOutputs(
+        shared_embedding=tf.ones((10, 10)),
+        control_logits=tf.ones((10, 1)),
+        treatment_logits=tf.ones((10, 1)),
+    )
+    _ = tfm.uplift.layers.encoders.concat_features.ConcatFeatures(['feature'])
+    _ = tfm.uplift.metrics.treatment_fraction.TreatmentFraction()
+    _ = tfm.uplift.losses.true_logits_loss.TrueLogitsLoss(tf_keras.losses.mse)
+
 
 if __name__ == '__main__':
   tf.test.main()

tensorflow_models/uplift/__init__.py

@@ -0,0 +1,23 @@
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""TensorFlow Models Uplift Libraries."""
+
+from official.recommendation.uplift import keys
+from official.recommendation.uplift import layers
+from official.recommendation.uplift import losses
+from official.recommendation.uplift import metrics
+from official.recommendation.uplift import models
+from official.recommendation.uplift import types
+from official.recommendation.uplift import utils

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

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

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

@@ -218,7 +218,7 @@ official/modeling/multitask/base_trainer_test.py,sha256=qJ7z4kid2XAX6hOIvUHa7dwq
 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=G4jwnWywn9tlEkdHgN8BApA3ymXFHIGnquq8rbVVBHo,4463
 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
@@ -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=qWMcLk1JmLkru5T9qcqO2UbB_AJQKjQ8uUSAxjR1l58,6883
-official/recommendation/uplift/metrics/loss_metric_test.py,sha256=QIUujZvUud-A2itI2vs2KpzWfLwbq2XDd9H9CFq5D4Y,14929
+official/recommendation/uplift/metrics/loss_metric.py,sha256=owN7A98TCc_UhOURvGfccaoVGOthdHdx1_fawEUGnmw,7289
+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/sliced_metric.py,sha256=O2I2apZK6IfOQK9Q_mgSiTUCnGokczp4e14zrrYNeRU,8564
 official/recommendation/uplift/metrics/sliced_metric_test.py,sha256=dhY41X8lqT_WW04XLjyjDerZwujEBGeTXtxf4NkYThw,11359
@@ -1199,13 +1199,14 @@ orbit/utils/summary_manager.py,sha256=oTdTQYDjH315rbJ__NFALldruHX5lLrwY34IU1KJwN
 orbit/utils/summary_manager_interface.py,sha256=MXXPd3aWqushByK-gh-Jk9A42VC5ETEEm7XY5DCB-9w,2253
 orbit/utils/tpu_summaries.py,sha256=vjfLJlNifFJayLo_nkdRD85XoisVBV2VDuzgZWZud6U,5641
 orbit/utils/tpu_summaries_test.py,sha256=ZRYguds3BcXn1gOXBHVJ5rGRZdeXQLa5w3IbwEBwGS4,4523
-tensorflow_models/__init__.py,sha256=etxw45SHxuwFCRX5qGxGMP83II0JfJulzNl5GSNJvhw,909
-tensorflow_models/tensorflow_models_test.py,sha256=AxUYUdiQn416UR7jg0h6rmv688esvlKDfpyDCIQkF18,1395
+tensorflow_models/__init__.py,sha256=n7s5ZGPDM-RqOQNldPyipjL8aEB4DdAQQUENYhgf_Q8,946
+tensorflow_models/tensorflow_models_test.py,sha256=nc6A9K53OGqF25xN5St8EiWvdVbdanoDB12hb32dzPc,1897
 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.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,,
+tf_models_nightly-2.17.0.dev20240411.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.17.0.dev20240411.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.17.0.dev20240411.dist-info/METADATA,sha256=I12qTeIZ3Y8oc91Y2TGdZ1WonCJu2KZaGuZdmIB7aa0,1432
+tf_models_nightly-2.17.0.dev20240411.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.17.0.dev20240411.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.17.0.dev20240411.dist-info/RECORD,,