diff --git a/examples/run_squad.py b/examples/run_squad.py index 028557b79f..6595f0464d 100644 --- a/examples/run_squad.py +++ b/examples/run_squad.py @@ -170,15 +170,19 @@ def train(args, train_dataset, model, tokenizer): steps_trained_in_current_epoch = 0 # Check if continuing training from a checkpoint if os.path.exists(args.model_name_or_path): - # set global_step to gobal_step of last saved checkpoint from model path - global_step = int(args.model_name_or_path.split("-")[-1].split("/")[0]) - epochs_trained = global_step // (len(train_dataloader) // args.gradient_accumulation_steps) - steps_trained_in_current_epoch = global_step % (len(train_dataloader) // args.gradient_accumulation_steps) + try: + # set global_step to gobal_step of last saved checkpoint from model path + checkpoint_suffix = args.model_name_or_path.split("-")[-1].split("/")[0] + global_step = int(checkpoint_suffix) + epochs_trained = global_step // (len(train_dataloader) // args.gradient_accumulation_steps) + steps_trained_in_current_epoch = global_step % (len(train_dataloader) // args.gradient_accumulation_steps) - logger.info(" Continuing training from checkpoint, will skip to saved global_step") - logger.info(" Continuing training from epoch %d", epochs_trained) - logger.info(" Continuing training from global step %d", global_step) - logger.info(" Will skip the first %d steps in the first epoch", steps_trained_in_current_epoch) + logger.info(" Continuing training from checkpoint, will skip to saved global_step") + logger.info(" Continuing training from epoch %d", epochs_trained) + logger.info(" Continuing training from global step %d", global_step) + logger.info(" Will skip the first %d steps in the first epoch", steps_trained_in_current_epoch) + except ValueError: + logger.info(" Starting fine-tuning.") tr_loss, logging_loss = 0.0, 0.0 model.zero_grad()