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