tfm.core.train_utils.BestCheckpointExporter

Keeps track of the best result, and saves its checkpoint.

tfm.core.train_utils.BestCheckpointExporter(
    export_dir: str, metric_name: str, metric_comp: str
)

Orbit will support an API for checkpoint exporter. This class will be used together with orbit once this functionality is ready.

export_dir The directory that will contain exported checkpoints.
metric_name Indicates which metric to look at, when determining which result is better. If eval_logs being passed to maybe_export_checkpoint is a nested dictionary, use | as a seperator for different layers.
metric_comp Indicates how to compare results. Either lower or higher.

best_ckpt_logs

best_ckpt_logs_path

best_ckpt_path Returns the best ckpt path or None if there is no ckpt yet.

Methods

export_best_eval_metric

View source

export_best_eval_metric(
    eval_logs, global_step
)

Export evaluation results of the best checkpoint into a json file.

maybe_export_checkpoint

View source

maybe_export_checkpoint(
    checkpoint, eval_logs, global_step, write_logs=True
) -> bool

Compare eval_logs with past eval_logs and export checkpoint if better.