def trace_model_from_checkpoint(logdir, logger, method_name='forward'): config_path = f'{logdir}/configs/_config.json' checkpoint_path = f'{logdir}/checkpoints/best.pth' logger.info('Load config') config = safitty.load(config_path) if 'distributed_params' in config: del config['distributed_params'] # Get expdir name # noinspection SpellCheckingInspection,PyTypeChecker # We will use copy of expdir from logs for reproducibility expdir_name = os.path.basename(config['args']['expdir']) expdir_from_logs = os.path.abspath(join(logdir, '../', expdir_name)) logger.info('Import experiment and runner from logdir') ExperimentType, RunnerType = \ import_experiment_and_runner(Path(expdir_from_logs)) experiment: Experiment = ExperimentType(config) logger.info('Load model state from checkpoints/best.pth') model = experiment.get_model(next(iter(experiment.stages))) checkpoint = utils.load_checkpoint(checkpoint_path) utils.unpack_checkpoint(checkpoint, model=model) logger.info('Tracing') traced = trace_model(model, experiment, RunnerType, method_name) logger.info('Done') return traced
def trace_model_from_checkpoint(logdir, method_name): config_path = logdir / "configs/_config.json" checkpoint_path = logdir / "checkpoints/best.pth" print("Load config") config: Dict[str, dict] = safitty.load_config(config_path) # Get expdir name config_expdir = Path(config["args"]["expdir"]) # We will use copy of expdir from logs for reproducibility expdir_from_logs = Path(logdir) / "code" / config_expdir.name print("Import experiment and runner from logdir") ExperimentType, RunnerType = \ import_experiment_and_runner(expdir_from_logs) experiment: Experiment = ExperimentType(config) print("Load model state from checkpoints/best.pth") model = experiment.get_model(next(iter(experiment.stages))) model.load_state_dict(torch.load(checkpoint_path)["model_state_dict"]) print("Tracing") traced = trace_model(model, experiment, RunnerType, method_name) print("Done") return traced
def trace_model_from_checkpoint(logdir, logger, method_name='forward', file='best'): config_path = f'{logdir}/configs/_config.json' checkpoint_path = f'{logdir}/checkpoints/{file}.pth' logger.info('Load config') config = safitty.load(config_path) if 'distributed_params' in config: del config['distributed_params'] # Get expdir name # noinspection SpellCheckingInspection,PyTypeChecker # We will use copy of expdir from logs for reproducibility expdir_name = config['args']['expdir'] logger.info(f'expdir_name from args: {expdir_name}') sys.path.insert(0, os.path.abspath(join(logdir, '../'))) expdir_from_logs = os.path.abspath(join(logdir, '../', expdir_name)) logger.info(f'expdir_from_logs: {expdir_from_logs}') logger.info('Import experiment and runner from logdir') ExperimentType, RunnerType = \ import_experiment_and_runner(Path(expdir_from_logs)) experiment: Experiment = ExperimentType(config) logger.info(f'Load model state from checkpoints/{file}.pth') model = experiment.get_model(next(iter(experiment.stages))) checkpoint = utils.load_checkpoint(checkpoint_path) utils.unpack_checkpoint(checkpoint, model=model) device = 'cpu' stage = list(experiment.stages)[0] loader = 0 mode = 'eval' requires_grad = False opt_level = None runner: RunnerType = RunnerType() runner.model, runner.device = model, device batch = experiment.get_native_batch(stage, loader) logger.info('Tracing') traced = trace_model( model, runner, batch, method_name=method_name, mode=mode, requires_grad=requires_grad, opt_level=opt_level, device=device, ) logger.info('Done') return traced
def trace_model_from_checkpoint( logdir: Path, method_name: str, checkpoint_name: str, mode: str = "eval", requires_grad: bool = False, ): config_path = logdir / "configs" / "_config.json" checkpoint_path = logdir / "checkpoints" / f"{checkpoint_name}.pth" print("Load config") config: Dict[str, dict] = safitty.load(config_path) # Get expdir name config_expdir = safitty.get(config, "args", "expdir", apply=Path) # We will use copy of expdir from logs for reproducibility expdir = Path(logdir) / "code" / config_expdir.name print("Import experiment and runner from logdir") ExperimentType, RunnerType = import_experiment_and_runner(expdir) experiment: Experiment = ExperimentType(config) print(f"Load model state from checkpoints/{checkpoint_name}.pth") model = experiment.get_model(next(iter(experiment.stages))) checkpoint = utils.load_checkpoint(checkpoint_path) utils.unpack_checkpoint(checkpoint, model=model) print("Tracing") traced = trace_model( model, experiment, RunnerType, method_name=method_name, mode=mode, requires_grad=requires_grad, ) print("Done") return traced
def trace( self, model: Model = None, batch=None, logdir: str = None, loader: DataLoader = None, method_name: str = "forward", mode: str = "eval", requires_grad: bool = False, fp16: Union[Dict, bool] = None, device: Device = "cpu", predict_params: dict = None, ) -> ScriptModule: """ Traces model using Torch Jit Args: model (Model): model to trace batch: batch to forward through the model to trace logdir (str, optional): If specified, the result will be written to the directory loader (DataLoader, optional): if batch is not specified, the batch will be ``next(iter(loader))`` method_name (str): model's method name that will be traced mode (str): ``train`` or ``eval`` requires_grad (bool): flag to trace with gradients fp16 (Union[Dict, bool]): If not None, then sets tracing params to FP16 deivice (Device): Torch deivice or a string predict_params (dict): additional parameters for model forward """ if batch is None: if loader is None: raise ValueError( "If batch is not provided the loader must be specified" ) batch = next(iter(loader)) if model is not None: self.model = model if isinstance(fp16, bool) and fp16: opt_level = "O1" elif isinstance(fp16, bool) and not fp16: opt_level = None elif isinstance(fp16, dict): opt_level = fp16["opt_level"] else: opt_level = fp16 if opt_level is not None: device = "cuda" elif device is None: if self.device is None: self.device = utils.get_device() device = self.device result = trace.trace_model( model=self.model, runner=self, batch=batch, method_name=method_name, mode=mode, requires_grad=requires_grad, opt_level=opt_level, device=device, predict_params=predict_params, ) if logdir is not None: filename = trace.get_trace_name( method_name=method_name, mode=mode, requires_grad=requires_grad, opt_level=opt_level, ) logdir = Path(logdir) output: Path = logdir / "trace" output.mkdir(exist_ok=True, parents=True) out_model = str(output / filename) torch.jit.save(result, out_model) return result
def trace_model_from_checkpoint( logdir: Path, method_name: str, checkpoint_name: str, stage: str = None, loader: Union[str, int] = None, mode: str = "eval", requires_grad: bool = False, opt_level: str = None, device: Device = "cpu", ): """Traces model using created experiment and runner. Args: logdir (Union[str, Path]): Path to Catalyst logdir with model checkpoint_name (str): Name of model checkpoint to use stage (str): experiment's stage name loader (Union[str, int]): experiment's loader name or its index method_name (str): Model's method name that will be used as entrypoint during tracing mode (str): Mode for model to trace (``train`` or ``eval``) requires_grad (bool): Flag to use grads opt_level (str): AMP FP16 init level device (str): Torch device Returns: the traced model """ config_path = logdir / "configs" / "_config.json" checkpoint_path = logdir / "checkpoints" / f"{checkpoint_name}.pth" print("Load config") config: Dict[str, dict] = utils.load_config(config_path) runner_params = config.get("runner_params", {}) or {} # Get expdir name config_expdir = Path(config["args"]["expdir"]) # We will use copy of expdir from logs for reproducibility expdir = Path(logdir) / "code" / config_expdir.name print("Import experiment and runner from logdir") ExperimentType, RunnerType = utils.import_experiment_and_runner(expdir) experiment: Experiment = ExperimentType(config) print(f"Load model state from checkpoints/{checkpoint_name}.pth") if stage is None: stage = list(experiment.stages)[0] model = experiment.get_model(stage) checkpoint = utils.load_checkpoint(checkpoint_path) utils.unpack_checkpoint(checkpoint, model=model) runner: RunnerType = RunnerType(**runner_params) runner.model, runner.device = model, device if loader is None: loader = 0 batch = experiment.get_native_batch(stage, loader) print("Tracing") traced = trace.trace_model( model=model, runner=runner, batch=batch, method_name=method_name, mode=mode, requires_grad=requires_grad, opt_level=opt_level, device=device, ) print("Done") return traced