tf-models-nightly 2.17.0.dev20240411__py2.py3-none-any.whl → 2.17.0.dev20240412__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.

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tf-models-nightly
-Version: 2.17.0.dev20240411
+Version: 2.17.0.dev20240412
 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.dev20240412.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
@@ -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.dev20240412.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.17.0.dev20240412.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.17.0.dev20240412.dist-info/METADATA,sha256=y_U6M920Hgob94pOpRJWfqistgFzPqBH0Ds0HOZOpoo,1432
+tf_models_nightly-2.17.0.dev20240412.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.17.0.dev20240412.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.17.0.dev20240412.dist-info/RECORD,,