def quantize_model_from_checkpoint( logdir: Path, checkpoint_name: str, stage: str = None, qconfig_spec: Optional[Union[Set, Dict]] = None, dtype: Optional[torch.dtype] = torch.qint8, backend: str = None, ) -> Model: """ Quantize 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 qconfig_spec: torch.quantization.quantize_dynamic parameter, you can define layers to be quantize dtype: type of the model parameters, default int8 backend: defines backend for quantization Returns: Quantized model """ if backend is not None: torch.backends.quantized.engine = backend config_path = logdir / "configs" / "_config.json" checkpoint_path = logdir / "checkpoints" / f"{checkpoint_name}.pth" logging.info("Load config") config: Dict[str, dict] = load_config(config_path) # 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 logger.info("Import experiment and runner from logdir") experiment: ConfigExperiment = None experiment, _, _ = prepare_config_api_components(expdir=expdir, config=config) logger.info(f"Load model state from checkpoints/{checkpoint_name}.pth") if stage is None: stage = list(experiment.stages)[0] model = experiment.get_model(stage) checkpoint = load_checkpoint(checkpoint_path) unpack_checkpoint(checkpoint, model=model) logger.info("Quantization is running...") quantized_model = quantization.quantize_dynamic( model.cpu(), qconfig_spec=qconfig_spec, dtype=dtype, ) logger.info("Done") return quantized_model
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" logging.info("Load config") config: Dict[str, dict] = load_config(config_path) # 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 logger.info("Import experiment and runner from logdir") experiment: ConfigExperiment = None experiment, runner, _ = prepare_config_api_components( expdir=expdir, config=config ) logger.info(f"Load model state from checkpoints/{checkpoint_name}.pth") if stage is None: stage = list(experiment.stages)[0] model = experiment.get_model(stage) checkpoint = load_checkpoint(checkpoint_path) unpack_checkpoint(checkpoint, model=model) runner.model, runner.device = model, device if loader is None: loader = 0 batch = get_native_batch_from_loaders( loaders=experiment.get_loaders(stage), loader=loader ) # function to run prediction on batch def predict_fn(model, inputs, **kwargs): # noqa: WPS442 model_dump = runner.model runner.model = model result = runner.predict_batch(inputs, **kwargs) runner.model = model_dump return result logger.info("Tracing is running...") traced_model = trace_model( model=model, predict_fn=predict_fn, batch=batch, method_name=method_name, mode=mode, requires_grad=requires_grad, opt_level=opt_level, device=device, ) logger.info("Done") return traced_model