docs

游雁

2023-05-25 b18f7d121f2f17df8bf2d0c2bbb223bc5ddbcc0f

 funasr/train/trainer.py

@@ -39,7 +39,7 @@
from funasr.torch_utils.device_funcs import to_device
from funasr.torch_utils.recursive_op import recursive_average
from funasr.torch_utils.set_all_random_seed import set_all_random_seed
from funasr.train.abs_espnet_model import AbsESPnetModel
from funasr.models.base_model import FunASRModel
from funasr.train.distributed_utils import DistributedOption
from funasr.train.reporter import Reporter
from funasr.train.reporter import SubReporter
@@ -143,11 +143,23 @@
        schedulers: Sequence[Optional[AbsScheduler]],
        scaler: Optional[GradScaler],
        ngpu: int = 0,
        oss_bucket=None,
    ):
        states = torch.load(
            checkpoint,
            map_location=f"cuda:{torch.cuda.current_device()}" if ngpu > 0 else "cpu",
        )
        if oss_bucket is None:
            if os.path.exists(checkpoint):
                states = torch.load(
                    checkpoint,
                    map_location=f"cuda:{torch.cuda.current_device()}" if ngpu > 0 else "cpu",
                )
            
            else:
                return 0
        else:
            if oss_bucket.object_exists(checkpoint):
                buffer = BytesIO(oss_bucket.get_object(checkpoint).read())
                states = torch.load(buffer, map_location=f"cuda:{torch.cuda.current_device()}" if ngpu > 0 else "cpu",)
            else:
                return 0
        model.load_state_dict(states["model"])
        reporter.load_state_dict(states["reporter"])
        for optimizer, state in zip(optimizers, states["optimizers"]):
@@ -166,7 +178,7 @@
    @classmethod
    def run(
        cls,
        model: AbsESPnetModel,
        model: FunASRModel,
        optimizers: Sequence[torch.optim.Optimizer],
        schedulers: Sequence[Optional[AbsScheduler]],
        train_iter_factory: AbsIterFactory,
@@ -206,15 +218,16 @@
        else:
            scaler = None

        if trainer_options.resume and (output_dir / "checkpoint.pb").exists():
        if trainer_options.resume:
            cls.resume(
                checkpoint=output_dir / "checkpoint.pb",
                checkpoint=os.path.join(trainer_options.output_dir, "checkpoint.pb") if trainer_options.use_pai else output_dir / "checkpoint.pb",
                model=model,
                optimizers=optimizers,
                schedulers=schedulers,
                reporter=reporter,
                scaler=scaler,
                ngpu=trainer_options.ngpu,
                oss_bucket=trainer_options.oss_bucket if trainer_options.use_pai else None,
            )

        start_epoch = reporter.get_epoch() + 1