Module: orbit.actions

Defines an "action" abstraction for use with orbit.Controller.

"Actions" are simply arbitrary callables that are applied by the Controller to the output of train steps (after each inner loop of steps_per_loop steps) or an evaluation. This provides a hook mechanism, enabling things like reporting metrics to Vizier, model exporting, additional logging, etc.

The basic Action abstraction (just a type alias) is defined in the controller module. This actions module adds a ConditionalAction utility class to make it easy to trigger actions conditionally based on reusable predicates, as well as a small handful of predefined conditions/actions (in particular, a NewBestMetric condition and an ExportSavedModel action).

One example of using actions to do metric-conditional export:

new_best_metric = orbit.actions.NewBestMetric('accuracy')
export_action = orbit.actions.ConditionalAction(
    condition=lambda x: x['accuracy'] > 0.9 and new_best_metric(x),
    action=orbit.actions.ExportSavedModel(
        model,
        orbit.actions.ExportFileManager(
            base_name=f'{FLAGS.model_dir}/saved_model',
            next_id_fn=trainer.global_step.numpy),
        signatures=model.infer))

controller = orbit.Controller(
    strategy=strategy,
    trainer=trainer,
    evaluator=evaluator,
    eval_actions=[export_action],
    global_step=trainer.global_step,
    steps_per_loop=FLAGS.steps_per_loop,
    checkpoint_manager=checkpoint_manager,
    summary_interval=1000)
client_0_actions = orbit.actions.ConditionalAction(
    condition=lambda _: client_id() == 0,
    action=[
        ...
    ])

In particular, the NewBestMetric condition may be used in multi-client settings if all clients are guaranteed to compute the same metric (ensuring this is up to client code, not Orbit). However, when saving metrics it may be helpful to avoid unnecessary writes by setting the write_value parameter to False for most clients.

Classes

class ConditionalAction: Represents an action that is only taken when a given condition is met.

class ExportFileManager: Utility class that manages a group of files with a shared base name.

class ExportSavedModel: Action that exports the given model as a SavedModel.

class JSONPersistedValue: Represents a value that is persisted via a file-based backing store.

class NewBestMetric: Condition that is satisfied when a new best metric is achieved.

class SaveCheckpointIfPreempted: Action that saves on-demand checkpoints after a preemption.