Module: tfm.core.train_utils

Training utils.

Classes

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

class ExperimentParser: Constructs the Experiment config from Flags or equivalent object.

class ParseConfigOptions: Use this dataclass instead of FLAGS to customize parse_configuration().

Functions

cast_leaf_nested_dict(...): Cast the leaves of a dictionary with arbitrary depth in place.

convert_variables_to_constants_v2_as_graph(...): Replaces all the variables in a graph with constants of the same values.

create_optimizer(...): A create optimizer util to be backward compatability with new args.

create_trainer(...): Create trainer.

get_leaf_nested_dict(...): Get leaf from a dictionary with arbitrary depth with a list of keys.

maybe_create_best_ckpt_exporter(...): Maybe create a BestCheckpointExporter object, according to the config.

parse_configuration(...): Parses ExperimentConfig from flags.

read_global_step_from_checkpoint(...): Read global step from checkpoint, or get global step from its filename.

remove_ckpts(...): Remove model checkpoints, so we can restart.

save_gin_config(...): Serializes and saves the experiment config.

serialize_config(...): Serializes and saves the experiment config.

try_count_flops(...): Counts and returns model FLOPs.

try_count_params(...): Count the number of parameters if model is possible.

write_json_summary(...): Dump evaluation metrics to json file.

write_model_params(...): Writes the model parameters and shapes to a file.

write_summary(...): Write evaluation metrics to TF summary.

BEST_CHECKPOINT_NAME 'best_ckpt'