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

official/modeling/multitask/interleaving_trainer.py

@@ -43,12 +43,6 @@ 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):
 
@@ -63,8 +57,6 @@ 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
 
@@ -73,6 +65,12 @@ 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.
@@ -99,6 +97,8 @@ 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/keys.py

@@ -17,7 +17,7 @@
 import enum
 
 
-class TwoTowerOutputKeys(enum.StrEnum):
+class TwoTowerOutputKeys(str, enum.Enum):
   """Keys for training and inference output tensors."""
 
   CONTROL_PREDICTIONS = "control_predictions"

official/recommendation/uplift/layers/heads/two_tower_logits_head.py

@@ -24,7 +24,7 @@ import tensorflow as tf, tf_keras
 
 
 @enum.unique
-class LayeringMethod(enum.StrEnum):
+class LayeringMethod(str, enum.Enum):
   """Layering method between the control and treatment towers."""
 
   # No layering.

official/recommendation/uplift/metrics/loss_metric.py

@@ -183,6 +183,9 @@ class LossMetric(tf_keras.metrics.Metric):
   def result(self) -> tf.Tensor | dict[str, tf.Tensor]:
     return self._loss.result()
 
+  def reset_state(self):
+    self._loss.reset_state()
+
   def get_config(self) -> dict[str, Any]:
     config = super().get_config()
     config["loss_fn"] = tf_keras.utils.serialize_keras_object(self._loss_fn)

official/recommendation/uplift/metrics/sliced_metric.py

@@ -193,6 +193,11 @@ class SlicedMetric(tf_keras.metrics.Metric):
         f"{metric_result}."
     )
 
+  def reset_state(self):
+    self._metric.reset_state()
+    for metric in self._sliced_metrics:
+      metric.reset_state()
+
   def get_config(self):
     return {
         "name": self.name,

official/recommendation/uplift/metrics/sliced_metric_test.py

@@ -315,6 +315,33 @@ class SlicedMetricTest(keras_test_case.KerasTestCase, parameterized.TestCase):
     }
     self.assertDictEqual(expected_result, metric.result())
 
+  def test_reset_state(self):
+    metric = sliced_metric.SlicedMetric(
+        metric=tf_keras.metrics.AUC(curve="PR", from_logits=False, name="auc"),
+        slicing_spec={"control": False, "treatment": True},
+    )
+
+    expected_initial_result = {
+        "auc": 0.0,
+        "auc/control": 0.0,
+        "auc/treatment": 0.0,
+    }
+    self.assertAllClose(expected_initial_result, metric.result())
+
+    metric.update_state(
+        tf.constant([[0], [0], [1], [1]]),  # y_true
+        tf.constant([[0.2], [0.6], [0.3], [0.7]]),  # y_pred
+        slicing_feature=tf.constant([[True], [False], [True], [False]]),
+    )
+
+    result = metric.result()
+    self.assertGreater(result["auc"], 0.0)
+    self.assertGreater(result["auc/control"], 0.0)
+    self.assertGreater(result["auc/treatment"], 0.0)
+
+    metric.reset_state()
+    self.assertAllClose(expected_initial_result, metric.result())
+
   def test_metric_config(self):
     metric = sliced_metric.SlicedMetric(
         tf_keras.metrics.SparseTopKCategoricalAccuracy(k=2, name="accuracy@2"),

official/recommendation/uplift/models/two_tower_uplift_model_test.py

@@ -21,6 +21,7 @@ from official.recommendation.uplift import keras_test_case
 from official.recommendation.uplift import keys
 from official.recommendation.uplift.layers.uplift_networks import two_tower_uplift_network
 from official.recommendation.uplift.losses import true_logits_loss
+from official.recommendation.uplift.metrics import loss_metric
 from official.recommendation.uplift.models import two_tower_uplift_model
 
 
@@ -127,6 +128,53 @@ class TwoTowerUpliftModelTest(
     }
     self.assertAllClose(expected_predictions, model.predict(dataset))
 
+  def test_classification_model_trains(self):
+    tf_keras.utils.set_random_seed(1)
+
+    # Create binary classifier uplift model.
+    uplift_network = self._get_uplift_network(
+        control_feature_encoder=None, control_input_combiner=None
+    )
+    model = two_tower_uplift_model.TwoTowerUpliftModel(
+        treatment_indicator_feature_name="is_treatment",
+        uplift_network=uplift_network,
+        inverse_link_fn=tf.math.sigmoid,
+    )
+    model.compile(
+        optimizer=tf_keras.optimizers.SGD(0.1),
+        loss=true_logits_loss.TrueLogitsLoss(
+            loss_fn=tf_keras.losses.binary_crossentropy, from_logits=True
+        ),
+        metrics=[
+            loss_metric.LossMetric(
+                tf_keras.metrics.AUC(curve="PR", from_logits=True, name="aucpr")
+            ),
+        ],
+    )
+
+    # Create toy classification dataset.
+    treatment = tf.constant([[1], [1], [0], [1], [1], [1], [0], [1], [0], [1]])
+    y = treatment
+    dataset = tf.data.Dataset.from_tensor_slices((
+        {
+            "shared_feature": np.random.normal(size=(10, 1)),
+            "treatment_feature": np.random.normal(size=(10, 1)),
+            "is_treatment": treatment,
+        },
+        y,
+    )).batch(5)
+
+    # Test model training.
+    history = model.fit(dataset, epochs=100)
+    self.assertIn("loss", history.history)
+    self.assertLen(history.history["loss"], 100)
+    self.assertBetween(
+        history.history["loss"][-1], 0.0, history.history["loss"][0]
+    )
+    self.assertIn("aucpr", history.history)
+    self.assertLess(history.history["aucpr"][0], 1.0)
+    self.assertEqual(history.history["aucpr"][-1], 1.0)
+
   @parameterized.named_parameters(
       {
           "testcase_name": "identity",

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

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

{tf_models_nightly-2.17.0.dev20240411.dist-info → tf_models_nightly-2.17.0.dev20240413.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=G4jwnWywn9tlEkdHgN8BApA3ymXFHIGnquq8rbVVBHo,4463
+official/modeling/multitask/interleaving_trainer.py,sha256=ZZHKsqbJKLqvwtgy-PUv_S_8bDG0MhJDNwWICY_IF6Q,4458
 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
@@ -887,7 +887,7 @@ official/recommendation/ranking/data/data_pipeline_multi_hot_test.py,sha256=arLj
 official/recommendation/ranking/data/data_pipeline_test.py,sha256=VRYo7WqURRkM3lbmfctvSZxyH1EzUfqwxR3sy2sZxdc,2345
 official/recommendation/uplift/__init__.py,sha256=_jZilTPWKu-MfMaz1IgBjEW6wqkK3FNZ1QAP4a8my3I,990
 official/recommendation/uplift/keras_test_case.py,sha256=gF5Z2FzXlKAvhuJDdj7PmFj3jsW_ZmUAv_F9Xokvs2M,6156
-official/recommendation/uplift/keys.py,sha256=jXrgCwSl1Ma4sThUeh8GdeQAig4lzdSbgrOwBcjvQSE,1030
+official/recommendation/uplift/keys.py,sha256=7zkxkPIcXceIN5hWm4ATai4h8ymwCj85dU2r8-3XifM,1032
 official/recommendation/uplift/types.py,sha256=OBpMvU4uAOKXGwN1UIPZw6KWJXBXeuZGT3P3VyDkyWc,4769
 official/recommendation/uplift/utils.py,sha256=ZSrzoFJosSdmu7P3yYDo-TER8UQKFimy9ntuXqV4L4c,3501
 official/recommendation/uplift/utils_test.py,sha256=K-pzTe3MoOEjdy5tLNSRpIQoWWH12MjFTPkhEIHfS3A,4610
@@ -896,7 +896,7 @@ official/recommendation/uplift/layers/encoders/__init__.py,sha256=QaQk0BMHJNE8Pc
 official/recommendation/uplift/layers/encoders/concat_features.py,sha256=UNjDXN2C-GBXJaiitMYmKuhIsOiR1ROfJ73vHkuGkoA,4210
 official/recommendation/uplift/layers/encoders/concat_features_test.py,sha256=sIjdoXqQHfQQ5P0yOz_PCuGQHl6j-wKNzaN42PL1jKI,9368
 official/recommendation/uplift/layers/heads/__init__.py,sha256=nbtNxVbIGZq4GLEw9APPwCX3Zhatv5qwc0ddvAaelvU,728
-official/recommendation/uplift/layers/heads/two_tower_logits_head.py,sha256=1HkR2LP3nBbAQrYcbGHsQPO36UuVQ8Pb8CAjWFT0c1Q,6613
+official/recommendation/uplift/layers/heads/two_tower_logits_head.py,sha256=EyeVdiRNe4qk1aEvYEp1tEeGdllUFqtt8saCrK_QI1M,6615
 official/recommendation/uplift/layers/heads/two_tower_logits_head_test.py,sha256=GKk8kYgq_gXcheUfcG1u_i5I5nVh8jRdam9KxvEHsXU,7615
 official/recommendation/uplift/layers/uplift_networks/__init__.py,sha256=eaM75bIO4WZRABzDjTl5bIsR5kKiICpzw-QpgXh7K0A,917
 official/recommendation/uplift/layers/uplift_networks/base_uplift_networks.py,sha256=pGVmdNXyDZR0TA9Htxa07KoUOzE0uEl5IyHEAVRgTk4,1329
@@ -912,11 +912,11 @@ 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=owN7A98TCc_UhOURvGfccaoVGOthdHdx1_fawEUGnmw,7289
+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/sliced_metric.py,sha256=O2I2apZK6IfOQK9Q_mgSiTUCnGokczp4e14zrrYNeRU,8564
-official/recommendation/uplift/metrics/sliced_metric_test.py,sha256=dhY41X8lqT_WW04XLjyjDerZwujEBGeTXtxf4NkYThw,11359
+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
 official/recommendation/uplift/metrics/treatment_fraction_test.py,sha256=LtFljDdz9yfH1GNDMo8OcdS4yhsez5WyHsthH3qJf3s,5430
 official/recommendation/uplift/metrics/treatment_sliced_metric.py,sha256=S0ZSoOHcjeWDWiEZlRnFHtRkOzizvrfmsFwbYP0Z0rY,3804
@@ -927,7 +927,7 @@ official/recommendation/uplift/metrics/variance.py,sha256=rhwZzUX-cRbwr-7vhC0I0b
 official/recommendation/uplift/metrics/variance_test.py,sha256=EPISeHOFIh6WfODuC0SXbnmMugh90acMmm4BJkEZXlo,7757
 official/recommendation/uplift/models/__init__.py,sha256=kWy2K5LGXHVyyrTjJvbVFcBjj1bjPRI2dpIq-sfdhvo,716
 official/recommendation/uplift/models/two_tower_uplift_model.py,sha256=Fb6nLFAOqch81ravK57K9kggAeqvtJcBtKGZwCex0ts,5028
-official/recommendation/uplift/models/two_tower_uplift_model_test.py,sha256=yvg-FMq66tqD9IYC0lQWZcKvs_bqmiKc5s8p2K2FFrw,8361
+official/recommendation/uplift/models/two_tower_uplift_model_test.py,sha256=J7qC9f0fDG1aIrLz85K1qUzTFyAIH0v8eA1yfPJb9YY,10061
 official/utils/__init__.py,sha256=7oiypy0N82PDw9aSdcJBLVoGTd_oRSUOdvuJhMv4leQ,609
 official/utils/hyperparams_flags.py,sha256=2FCAxfblio6ay36Yf4o7Nx188wRzFM1mbKOtVXiZCzo,4607
 official/utils/docs/__init__.py,sha256=7oiypy0N82PDw9aSdcJBLVoGTd_oRSUOdvuJhMv4leQ,609
@@ -1204,9 +1204,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.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,,
+tf_models_nightly-2.17.0.dev20240413.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.17.0.dev20240413.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.17.0.dev20240413.dist-info/METADATA,sha256=DqjFt5jNaqGygdbDJo34myYN5F7c6EboWUtItI2AQVQ,1432
+tf_models_nightly-2.17.0.dev20240413.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.17.0.dev20240413.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.17.0.dev20240413.dist-info/RECORD,,