| | |
| | | |
| | | from contextlib import nullcontext |
| | | import torch.distributed as dist |
| | | from collections.abc import Sequence |
| | | |
| | | from omegaconf import DictConfig, OmegaConf |
| | | from torch.cuda.amp import autocast, GradScaler |
| | | from torch.nn.parallel import DistributedDataParallel as DDP |
| | |
| | | from funasr.train_utils.trainer import Trainer |
| | | from funasr.schedulers import scheduler_classes |
| | | from funasr.train_utils.initialize import initialize |
| | | from funasr.download.download_from_hub import download_model |
| | | from funasr.download.download_model_from_hub import download_model |
| | | from funasr.models.lora.utils import mark_only_lora_as_trainable |
| | | from funasr.train_utils.set_all_random_seed import set_all_random_seed |
| | | from funasr.train_utils.load_pretrained_model import load_pretrained_model |
| | |
| | | if freeze_param is not None: |
| | | if "," in freeze_param: |
| | | freeze_param = eval(freeze_param) |
| | | if not isinstance(freeze_param, Sequence): |
| | | if not isinstance(freeze_param, (list, tuple)): |
| | | freeze_param = (freeze_param,) |
| | | logging.info("freeze_param is not None: %s", freeze_param) |
| | | for t in freeze_param: |
| | |
| | | try: |
| | | from tensorboardX import SummaryWriter |
| | | |
| | | writer = SummaryWriter(tensorboard_dir) if trainer.rank == 0 else None |
| | | writer = SummaryWriter(tensorboard_dir) # if trainer.rank == 0 else None |
| | | except: |
| | | writer = None |
| | | |
| | | dataloader_tr, dataloader_val = None, None |
| | | for epoch in range(trainer.start_epoch, trainer.max_epoch + 1): |
| | | for epoch in range(trainer.start_epoch, trainer.max_epoch): |
| | | time1 = time.perf_counter() |
| | | |
| | | for data_split_i in range(trainer.start_data_split_i, dataloader.data_split_num): |
| | | time_slice_i = time.perf_counter() |
| | | dataloader_tr, dataloader_val = dataloader.build_iter( |
| | | epoch, data_split_i=data_split_i, start_step=trainer.start_step |
| | | ) |
| | | trainer.start_step = 0 |
| | | |
| | | trainer.train_epoch( |
| | | model=model, |
| | | optim=optim, |
| | |
| | | writer=writer, |
| | | data_split_i=data_split_i, |
| | | data_split_num=dataloader.data_split_num, |
| | | start_step=trainer.start_step, |
| | | ) |
| | | trainer.start_step = 0 |
| | | |
| | | torch.cuda.empty_cache() |
| | | |
| | | time_escaped = (time.perf_counter() - time_slice_i) / 3600.0 |
| | | logging.info( |
| | | f"rank: {local_rank}, " |
| | | f"time_escaped_epoch: {time_escaped:.3f} hours, " |
| | | f"estimated to finish {dataloader.data_split_num} data_slices, remaining: {dataloader.data_split_num-data_split_i} slices, {(dataloader.data_split_num-data_split_i)*time_escaped:.3f} hours, " |
| | | f"epoch: {trainer.max_epoch - epoch} epochs, {((trainer.max_epoch - epoch - 1)*dataloader.data_split_num + dataloader.data_split_num-data_split_i)*time_escaped:.3f} hours\n" |
| | | ) |
| | | |
| | | trainer.start_data_split_i = 0 |
| | | trainer.validate_epoch( |
| | | model=model, dataloader_val=dataloader_val, epoch=epoch, writer=writer |
| | | model=model, dataloader_val=dataloader_val, epoch=epoch + 1, writer=writer |
| | | ) |
| | | scheduler.step() |
| | | |
| | | trainer.save_checkpoint(epoch, model=model, optim=optim, scheduler=scheduler, scaler=scaler) |
| | | trainer.step_in_epoch = 0 |
| | | trainer.save_checkpoint( |
| | | epoch + 1, model=model, optim=optim, scheduler=scheduler, scaler=scaler |
| | | ) |
| | | |
| | | time2 = time.perf_counter() |
| | | time_escaped = (time2 - time1) / 3600.0 |
| | |
| | | f"estimated to finish {trainer.max_epoch} " |
| | | f"epoch: {(trainer.max_epoch - epoch) * time_escaped:.3f} hours\n" |
| | | ) |
| | | trainer.train_acc_avg = 0.0 |
| | | trainer.train_loss_avg = 0.0 |
| | | |
| | | if trainer.rank == 0: |
| | | average_checkpoints(trainer.output_dir, trainer.avg_nbest_model) |