| | |
| | | from funasr.train_utils.recursive_op import recursive_average |
| | | from funasr.train_utils.average_nbest_models import average_checkpoints |
| | | from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler |
| | | import funasr.utils.misc as misc_utils |
| | | |
| | | try: |
| | | import wandb |
| | |
| | | Args: |
| | | epoch (int): The epoch number at which the checkpoint is being saved. |
| | | """ |
| | | if self.use_ddp or self.use_fsdp: |
| | | dist.barrier() |
| | | step_in_epoch = None if step is None else step_in_epoch |
| | | if self.use_deepspeed: |
| | | |
| | |
| | | filename = os.path.join(self.output_dir, key) |
| | | logging.info(f"Delete: {filename}") |
| | | if os.path.exists(filename): |
| | | os.remove(filename) |
| | | # os.remove(filename) |
| | | misc_utils.smart_remove(filename) |
| | | |
| | | elif self.use_fsdp: |
| | | pass |
| | |
| | | filename = os.path.join(self.output_dir, key) |
| | | logging.info(f"Delete: {filename}") |
| | | if os.path.exists(filename): |
| | | os.remove(filename) |
| | | # os.remove(filename) |
| | | misc_utils.smart_remove(filename) |
| | | |
| | | if self.use_ddp or self.use_fsdp: |
| | | dist.barrier() |
| | |
| | | ckpt = os.path.join(self.output_dir, "model.pt") |
| | | if os.path.exists(ckpt): |
| | | _, checkpoint = model.load_checkpoint(self.output_dir, "model.pt") |
| | | |
| | | self.start_epoch = checkpoint["epoch"] |
| | | self.saved_ckpts = checkpoint["saved_ckpts"] |
| | | self.val_acc_step_or_eoch = ( |
| | | checkpoint["val_acc_step_or_eoch"] |
| | |
| | | self.train_acc_avg = train_acc_avg.detach().cpu().item() / self.world_size |
| | | |
| | | def forward_step(self, model, batch, loss_dict={}): |
| | | dtype = torch.bfloat16 |
| | | with maybe_autocast(dtype=self.dtype, use_deepspeed=self.use_deepspeed): |
| | | retval = model(**batch) |
| | | |
| | |
| | | "data_split_i": kwargs.get("data_split_i", 0), |
| | | "data_split_num": kwargs.get("data_split_num", 1), |
| | | "log_step": batch_idx + kwargs.get("start_step", 0), |
| | | "batch_total": batch_idx, |
| | | "step_in_epoch": batch_idx, |
| | | "batch_total": batch_idx + 1, |
| | | "step_in_epoch": batch_idx + 1, |
| | | "lr": 0.0, |
| | | } |
| | | |
| | |
| | | ckpt_name = f'model.pt.ep{epoch}.{kwargs.get("step_in_epoch")}' |
| | | self.val_acc_step_or_eoch[ckpt_name] = self.val_acc_avg |
| | | self.val_loss_step_or_eoch[ckpt_name] = self.val_loss_avg |
| | | |
| | | if self.use_ddp or self.use_fsdp or self.use_deepspeed: |
| | | dist.barrier() |
| | | |
| | | model.train() |
| | | |
| | | def log( |