def trace(
self,
*,
model: Model = None,
batch: Any = 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
device (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
assert self.model is not None
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 = utils.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 = utils.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(
self,
*,
model: Model = None,
batch: Any = 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
device (Device): Torch device 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
assert self.model is not None
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
# Dumping previous state of the model, we will need it to restore
_device, _is_training, _requires_grad = (
self.device,
self.model.training,
utils.get_requires_grad(self.model),
)
self.model.to(device)
# function to run prediction on batch
def predict_fn(model, inputs, **kwargs):
_model = self.model
self.model = model
result = self.predict_batch(inputs, **kwargs)
self.model = _model
return result
traced_model = utils.trace_model(
model=self.model,
predict_fn=predict_fn,
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:
utils.save_traced_model(
model=traced_model,
logdir=logdir,
method_name=method_name,
mode=mode,
requires_grad=requires_grad,
opt_level=opt_level,
)
# Restore previous state of the model
getattr(self.model, "train" if _is_training else "eval")()
utils.set_requires_grad(self.model, _requires_grad)
self.model.to(_device)
return traced_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"
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 = utils.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
