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

official/modeling/multitask/train_lib.py

@@ -78,15 +78,8 @@ def run_experiment(
   is_training = 'train' in mode
   is_eval = 'eval' in mode
   with distribution_strategy.scope():
-    optimizer = train_utils.create_optimizer(task, params)
-    kwargs = dict(multi_task=task, multi_task_model=model, optimizer=optimizer)
-    if params.trainer.trainer_type == 'interleaving':
-      sampler = task_sampler.get_task_sampler(params.trainer.task_sampler,
-                                              task.task_weights)
-      kwargs.update(dict(task_sampler=sampler))
-    if trainer is None:
-      trainer = TRAINERS[params.trainer.trainer_type](
-          **kwargs) if is_training else None
+    if is_training and trainer is None:
+      trainer = get_trainer(distribution_strategy, params, task, model)
     if is_eval:
       eval_steps = task.task_eval_steps
       evaluator = evaluator_lib.MultiTaskEvaluator(
@@ -157,6 +150,38 @@ def run_experiment(
       return model
 
 
+def get_trainer(
+    distribution_strategy: tf.distribute.Strategy,
+    params: configs.MultiEvalExperimentConfig,
+    task: multitask.MultiTask,
+    model: base_model.MultiTaskBaseModel | tf_keras.Model,
+) -> orbit.StandardTrainer:
+  """Creates a multi-task trainer for the given task.
+
+  Args:
+    distribution_strategy: A distribution strategy.
+    params: ExperimentConfig instance.
+    task: A MultiTaskTask instance.
+    model: A MultiTaskBaseModel instance.
+
+  Returns:
+      An Orbit trainer instance.
+  """
+  with distribution_strategy.scope():
+    kwargs = dict(
+        multi_task=task,
+        multi_task_model=model,
+        optimizer=train_utils.create_optimizer(task, params),
+    )
+    if params.trainer.trainer_type == 'interleaving':
+      kwargs.update(
+          task_sampler=task_sampler.get_task_sampler(
+              params.trainer.task_sampler, task.task_weights
+          )
+      )
+    return TRAINERS[params.trainer.trainer_type](**kwargs)
+
+
 TrainActionsFactoryType = Callable[
     [
         configs.MultiEvalExperimentConfig,

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

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

{tf_models_nightly-2.19.0.dev20250226.dist-info → tf_models_nightly-2.19.0.dev20250227.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=ZHRN6afigqcupAz61P6qBEh_RhomAwVgvjgkRHSv0eU,12827
+official/modeling/multitask/train_lib.py,sha256=TbqjoGlrNxxbq3XxKZ04PTxtZFQUlfwX5sd1-9dm5QI,13403
 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.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,,
+tf_models_nightly-2.19.0.dev20250227.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.19.0.dev20250227.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.19.0.dev20250227.dist-info/METADATA,sha256=c2C7Tg_sj_ybjTh9TtaYbTlf7RMk16JqrLmVI6K3DUI,1432
+tf_models_nightly-2.19.0.dev20250227.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.19.0.dev20250227.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.19.0.dev20250227.dist-info/RECORD,,