update

speech_asr

2023-04-18 a1a79bbe3e971a00bc315d011a2e0764b3bc3111

 funasr/train/trainer.py

@@ -94,7 +94,7 @@
    wandb_model_log_interval: int
    use_pai: bool
    oss_bucket: Union[oss2.Bucket, None]

    batch_interval: int

class Trainer:
    """Trainer having a optimizer.
@@ -186,7 +186,10 @@
                logging.warning("No keep_nbest_models is given. Change to [1]")
                trainer_options.keep_nbest_models = [1]
            keep_nbest_models = trainer_options.keep_nbest_models

     
        #assert batch_interval is set and >0
        assert trainer_options.batch_interval > 0
 
        output_dir = Path(trainer_options.output_dir)
        reporter = Reporter()
        if trainer_options.use_amp:
@@ -560,12 +563,33 @@
        # [For distributed] Because iteration counts are not always equals between
        # processes, send stop-flag to the other processes if iterator is finished
        iterator_stop = torch.tensor(0).to("cuda" if ngpu > 0 else "cpu")

        
        #get the rank
        rank = distributed_option.dist_rank
        #get the num batch updates
        num_batch_updates = 0
        #ouput dir
        output_dir = Path(options.output_dir)
        #batch interval
        batch_interval = options.batch_interval       
        assert batch_interval > 0
 
        start_time = time.perf_counter()
        for iiter, (_, batch) in enumerate(
            reporter.measure_iter_time(iterator, "iter_time"), 1
        ):
            assert isinstance(batch, dict), type(batch)

            if rank == 0:
                if hasattr(model, "num_updates") or (hasattr(model, "module") and hasattr(model.module, "num_updates")):
                    num_batch_updates = model.get_num_updates() if hasattr(model,"num_updates") else model.module.get_num_updates()
                if (num_batch_updates%batch_interval == 0) and (options.oss_bucket is not None):
                    if options.use_pai:
                        buffer = BytesIO()
                        torch.save(model.state_dict(), buffer)
                        options.oss_bucket.put_object(os.path.join(output_dir, f"{num_batch_updates}step.pb"), buffer.getvalue())
                    else:
                        torch.save(model.state_dict(), os.path.join(output_dir, f"{num_batch_updates}step.pb"))

            if distributed:
                torch.distributed.all_reduce(iterator_stop, ReduceOp.SUM)
@@ -811,4 +835,4 @@
        else:
            if distributed:
                iterator_stop.fill_(1)
                torch.distributed.all_reduce(iterator_stop, ReduceOp.SUM)
                torch.distributed.all_reduce(iterator_stop, ReduceOp.SUM)