Runs experiment with Orbit training loop. The default experiment runner for model garden experiments. User can customize the experiment pipeline by subclassing this class and replacing components or functions. For example, an experiment runner with customized checkpoint manager: ```pyth
| 32 | |
| 33 | |
| 34 | class OrbitExperimentRunner: |
| 35 | """Runs experiment with Orbit training loop. |
| 36 | |
| 37 | The default experiment runner for model garden experiments. User can |
| 38 | customize the experiment pipeline by subclassing this class and replacing |
| 39 | components or functions. |
| 40 | |
| 41 | For example, an experiment runner with customized checkpoint manager: |
| 42 | |
| 43 | ```python |
| 44 | class MyExpRunnerWithExporter(OrbitExperimentRunner): |
| 45 | def _maybe_build_checkpoint_manager(sefl): |
| 46 | # Replaces the default CheckpointManger with a customized one. |
| 47 | return MyCheckpointManager(*args) |
| 48 | |
| 49 | # In user code, instead of the orginal |
| 50 | # `OrbitExperimentRunner(..).run(mode)`, now user can do: |
| 51 | MyExpRunnerWithExporter(**needed_kwargs).run(mode) |
| 52 | ``` |
| 53 | |
| 54 | Similar override can be done to other components. |
| 55 | """ |
| 56 | |
| 57 | def __init__( |
| 58 | self, |
| 59 | distribution_strategy: tf.distribute.Strategy, |
| 60 | task: base_task.Task, |
| 61 | mode: str, |
| 62 | params: config_definitions.ExperimentConfig, |
| 63 | model_dir: str, |
| 64 | run_post_eval: bool = False, |
| 65 | save_summary: bool = True, |
| 66 | train_actions: Optional[List[orbit.Action]] = None, |
| 67 | eval_actions: Optional[List[orbit.Action]] = None, |
| 68 | trainer: Optional[base_trainer.Trainer] = None, |
| 69 | controller_cls=orbit.Controller, |
| 70 | summary_manager: Optional[orbit.utils.SummaryManager] = None, |
| 71 | eval_summary_manager: Optional[orbit.utils.SummaryManager] = None, |
| 72 | enable_async_checkpointing: bool = False, |
| 73 | ): |
| 74 | """Constructor. |
| 75 | |
| 76 | Args: |
| 77 | distribution_strategy: A distribution strategy. |
| 78 | task: A Task instance. |
| 79 | mode: A 'str', specifying the mode. Can be 'train', 'eval', |
| 80 | 'train_and_eval' or 'continuous_eval'. |
| 81 | params: ExperimentConfig instance. |
| 82 | model_dir: A 'str', a path to store model checkpoints and summaries. |
| 83 | run_post_eval: Whether to run post eval once after training, metrics logs |
| 84 | are returned. |
| 85 | save_summary: Whether to save train and validation summary. |
| 86 | train_actions: Optional list of Orbit train actions. |
| 87 | eval_actions: Optional list of Orbit eval actions. |
| 88 | trainer: the base_trainer.Trainer instance. It should be created within |
| 89 | the strategy.scope(). |
| 90 | controller_cls: The controller class to manage the train and eval process. |
| 91 | Must be a orbit.Controller subclass. |