def load_existing_checkpoint(
checkpoint_path,
trainer,
restore_state=True) -> Tuple[bool, Optional[Dict]]:
loaded = False
extra_state = None
if not PathManager.isfile(checkpoint_path):
print(f"| No existing checkpoint at {checkpoint_path}. "
f"Starting training from scratch.")
return loaded, extra_state
if restore_state:
extra_state = trainer.load_checkpoint(checkpoint_path)
if extra_state is None:
loaded = False
print(
f"| Failed to load checkpoint and state from {checkpoint_path}."
)
else:
loaded = True
print(f"| Loaded checkpoint {checkpoint_path} "
f"(epoch {extra_state['epoch']}) with restored extra state.")
# batch_offset being None denotes this was a checkpoint saved at
# the end of an epoch (after the last batch).
if extra_state["batch_offset"] is None:
trainer.lr_step(extra_state["epoch"])
extra_state["epoch"] += 1
extra_state["batch_offset"] = 0
else:
dummy_state = trainer.load_checkpoint(checkpoint_path,
reset_optimizer=True)
if dummy_state is None:
loaded = False
print(
f"| Failed to load checkpoint weights from {checkpoint_path}.")
else:
loaded = True
print(f"| Loaded checkpoint weights from {checkpoint_path}.")
return loaded, extra_state
def setup_training_state(args, trainer, task, epoch_itr):
"""Set up the directory for saving checkpoints.
Load pretrained model if specified."""
PathManager.mkdirs(args.save_dir)
# If --restore-file is already present under --save-dir, use that one
# instead of --pretrained-checkpoint-file. The idea is that
# --pretrained-checkpoint-file allows the user to specify restoring from a
# different run's checkpoint (possibly with different training params),
# while not polluting the previous run's checkpoint directory
# with new checkpoints. However, if training gets interrupted
# and the user restarts training, we want to resume from
# the checkpoints under --save-dir, instead of
# restarting again from the old run's checkpoint at
# --pretrained-checkpoint-file.
#
# Note that if args.restore_file is an absolute path, os.path.join() will
# ignore previous directory args and just use the absolute path as is.
checkpoint_path = os.path.join(args.save_dir, args.restore_file)
restore_state = True
if PathManager.isfile(checkpoint_path):
print(
f"| Using --save-dir={args.save_dir}, --restore-file={args.restore_file}."
)
elif args.pretrained_checkpoint_file and PathManager.isfile(
args.pretrained_checkpoint_file):
checkpoint_path = args.pretrained_checkpoint_file
restore_state = args.load_pretrained_checkpoint_state
print(
f"| Using --pretrained-checkpoint-file={args.pretrained_checkpoint_file}, "
f"--load-pretrained-checkpoint-state={args.load_pretrained_checkpoint_state}."
)
extra_state = default_extra_state(args)
if not PathManager.isfile(
checkpoint_path) and args.multi_model_restore_files:
print(
f"| Restoring individual models from {args.multi_model_restore_files}"
)
multi_model.import_individual_models(args.multi_model_restore_files,
trainer)
else:
loaded, loaded_extra_state = checkpoint.load_existing_checkpoint(
checkpoint_path=checkpoint_path,
trainer=trainer,
restore_state=restore_state,
)
if loaded_extra_state:
extra_state.update(loaded_extra_state)
# Reset the start time for the current training run.
extra_state["start_time"] = time.time()
# Skips printing all training progress to prevent log spam.
training_progress = extra_state["training_progress"]
extra_state["training_progress"] = ([
"...truncated...", training_progress[-1]
] if len(training_progress) > 0 else [])
print(f"| extra_state: {extra_state}")
extra_state["training_progress"] = training_progress
epoch = extra_state["epoch"]
if extra_state["batch_offset"] == 0:
epoch -= 1 # this will be incremented when we call epoch_itr.next_epoch_itr()
epoch_itr.load_state_dict({
"epoch":
epoch,
"iterations_in_epoch":
extra_state["batch_offset"]
})
checkpoint_manager = None
if distributed_utils.is_master(args):
checkpoint_manager = checkpoint.CheckpointManager(
num_avg_checkpoints=args.num_avg_checkpoints,
auto_clear_checkpoints=args.auto_clear_checkpoints,
log_verbose=args.log_verbose,
checkpoint_files=extra_state["checkpoint_files"],
)
return extra_state, epoch_itr, checkpoint_manager
