tf-models-nightly 2.19.0.dev20250224__py2.py3-none-any.whl → 2.19.0.dev20250226__py2.py3-none-any.whl

official/modeling/multitask/train_lib.py

@@ -15,7 +15,7 @@
 """Multitask training driver library."""
 # pytype: disable=attribute-error
 import os
-from typing import Any, List, Mapping, Optional, Tuple, Union
+from typing import Any, List, Mapping, Optional, Tuple, Union, Callable
 from absl import logging
 import orbit
 import tensorflow as tf, tf_keras
@@ -157,6 +157,25 @@ def run_experiment(
       return model
 
 
+TrainActionsFactoryType = Callable[
+    [
+        configs.MultiEvalExperimentConfig,
+        orbit.StandardTrainer,
+        str,
+        tf.train.CheckpointManager,
+    ],
+    List[orbit.Action],
+]
+EvalActionsFactoryType = Callable[
+    [
+        configs.MultiEvalExperimentConfig,
+        orbit.AbstractEvaluator,
+        str,
+    ],
+    List[orbit.Action],
+]
+
+
 def run_experiment_with_multitask_eval(
     *,
     distribution_strategy: tf.distribute.Strategy,
@@ -171,6 +190,8 @@ def run_experiment_with_multitask_eval(
     eval_summary_manager: Optional[orbit.utils.SummaryManagerInterface] = None,
     best_ckpt_exporter_creator: Optional[Any] = train_utils
     .maybe_create_best_ckpt_exporter,
+    train_actions_factory: Optional[TrainActionsFactoryType] = None,
+    eval_actions_factory: Optional[EvalActionsFactoryType] = None,
 ) -> Tuple[Any, Any]:
   """Runs train/eval configured by the experiment params.
 
@@ -193,6 +214,8 @@ def run_experiment_with_multitask_eval(
       will be created internally for TensorBoard summaries by default from the
       `eval_summary_dir`.
     best_ckpt_exporter_creator: A functor for creating best checkpoint exporter.
+    train_actions_factory: Optional factory function to create train actions.
+    eval_actions_factory: Optional factory function to create eval actions.
 
   Returns:
       model: `tf_keras.Model` instance.
@@ -214,7 +237,6 @@ def run_experiment_with_multitask_eval(
 
     # Build the model or fetch the pre-cached one (which could be either
     # multi-task model or single task model).
-    model = None
     if trainer is None:
       if isinstance(train_task, multitask.MultiTask):
         model = train_task.build_multitask_model()
@@ -254,6 +276,23 @@ def run_experiment_with_multitask_eval(
       checkpoint_interval=params.trainer.checkpoint_interval,
       init_fn=trainer.initialize if trainer else None)
 
+  if trainer and train_actions_factory:
+    # pytype: disable=wrong-keyword-args
+    train_actions = train_actions_factory(
+        params=params,
+        trainer=trainer,
+        model_dir=model_dir,
+        checkpoint_manager=checkpoint_manager,
+    )
+    # pytype: enable=wrong-keyword-args
+  else:
+    train_actions = None
+
+  if evaluator and eval_actions_factory:
+    eval_actions = eval_actions_factory(params, evaluator, model_dir)
+  else:
+    eval_actions = None
+
   controller = orbit.Controller(
       strategy=distribution_strategy,
       trainer=trainer,
@@ -266,7 +305,10 @@ def run_experiment_with_multitask_eval(
       (save_summary) else None,
       eval_summary_manager=eval_summary_manager,
       summary_interval=params.trainer.summary_interval if
-      (save_summary) else None)
+      (save_summary) else None,
+      train_actions=train_actions,
+      eval_actions=eval_actions,
+      )
 
   logging.info('Starts to execute mode: %s', mode)
   with distribution_strategy.scope():

{tf_models_nightly-2.19.0.dev20250224.dist-info → tf_models_nightly-2.19.0.dev20250226.dist-info}/METADATA

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

{tf_models_nightly-2.19.0.dev20250224.dist-info → tf_models_nightly-2.19.0.dev20250226.dist-info}/RECORD

@@ -224,7 +224,7 @@ official/modeling/multitask/multitask.py,sha256=DV-ysfhPiIZgsrzZNylsPBxKNBf_xzPx
 official/modeling/multitask/task_sampler.py,sha256=SGVVdjMb5oG4vnCczpfdgBtbsdsXiyBLl9si_0V6nko,4897
 official/modeling/multitask/task_sampler_test.py,sha256=wkPTp1LCNx4uJbfHfVbKrQULQxumZS3ctBemEbnMokk,3037
 official/modeling/multitask/test_utils.py,sha256=fPi_TxtzHy_NGTYmdDAQay9TRjwG70UWHwysDqol4jw,4315
-official/modeling/multitask/train_lib.py,sha256=XmWRj-IOKijIo9EDCRQNRU9rj3qiIDhExiPxDN6mi6A,11633
+official/modeling/multitask/train_lib.py,sha256=ZHRN6afigqcupAz61P6qBEh_RhomAwVgvjgkRHSv0eU,12827
 official/modeling/multitask/train_lib_test.py,sha256=F_uDEQhSlXq9BH1QBlamUMQ_vRUflli7ef2P9MUKXeg,4766
 official/modeling/optimization/__init__.py,sha256=BI86b89P0xOksBlzOmbAdWOhlJQe6cxtGIRN_zqCwa8,1201
 official/modeling/optimization/adafactor_optimizer.py,sha256=zEHHFg9iH1UwcBSzzYE8bCNWln7QIIftXJEXTGTtCtE,792
@@ -1248,9 +1248,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.19.0.dev20250224.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.19.0.dev20250224.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.19.0.dev20250224.dist-info/METADATA,sha256=Cd694VFHARcVs9nk3pbtQAcpd7fP1TMWRRA265kwmkg,1432
-tf_models_nightly-2.19.0.dev20250224.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.19.0.dev20250224.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.19.0.dev20250224.dist-info/RECORD,,
+tf_models_nightly-2.19.0.dev20250226.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.19.0.dev20250226.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.19.0.dev20250226.dist-info/METADATA,sha256=pXPpfanhACUR5sKJVyUH2c8SY0vi-8tcXz-JPP71KpQ,1432
+tf_models_nightly-2.19.0.dev20250226.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.19.0.dev20250226.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.19.0.dev20250226.dist-info/RECORD,,