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_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)))
checkpoint = utils.load_checkpoint(checkpoint_path)
utils.unpack_checkpoint(checkpoint, model=model)
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 = 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(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 main(args, unknown_args):
args, config = parse_args_uargs(args, unknown_args)
set_global_seed(args.seed)
Experiment, Runner = import_experiment_and_runner(Path(args.expdir))
experiment = Experiment(config)
runner = Runner()
if experiment.logdir is not None:
dump_config(config, experiment.logdir, args.configs)
dump_code(args.expdir, experiment.logdir)
runner.run_experiment(experiment, check=args.check)
def main(args, unknown_args):
args, config = parse_args_uargs(args, unknown_args)
set_global_seed(args.seed)
prepare_cudnn(args.deterministic, args.benchmark)
Experiment, Runner = import_experiment_and_runner(Path(args.expdir))
experiment = Experiment(config)
runner = Runner()
if experiment.logdir is not None:
dump_environment(config, experiment.logdir, args.configs)
dump_code(args.expdir, experiment.logdir)
runner.run_experiment(experiment, check=args.check)
def main(args, unknown_args):
"""Run the ``catalyst-dl run`` script"""
args, config = parse_args_uargs(args, unknown_args)
set_global_seed(args.seed)
prepare_cudnn(args.deterministic, args.benchmark)
Experiment, Runner = import_experiment_and_runner(Path(args.expdir))
runner_params = config.pop("runner_params", {}) or {}
experiment = Experiment(config)
runner = Runner(**runner_params)
if experiment.logdir is not None:
dump_environment(config, experiment.logdir, args.configs)
dump_code(args.expdir, experiment.logdir)
check_run = safitty.get(config, "args", "check", default=False)
runner.run_experiment(experiment, check=check_run)
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 work(self):
args, config = self.parse_args_uargs()
set_global_seed(args.seed)
Experiment, R = import_experiment_and_runner(Path(args.expdir))
runner_params = config.pop('runner_params', {})
experiment = Experiment(config)
runner: Runner = R(**runner_params)
register()
self.experiment = experiment
self.runner = runner
stages = experiment.stages[:]
if self.master:
task = self.task if not self.task.parent \
else self.task_provider.by_id(self.task.parent)
task.steps = len(stages)
self.task_provider.commit()
self._checkpoint_fix_config(experiment)
_get_callbacks = experiment.get_callbacks
def get_callbacks(stage):
res = self.callbacks()
for k, v in _get_callbacks(stage).items():
res[k] = v
self._checkpoint_fix_callback(res)
return res
experiment.get_callbacks = get_callbacks
if experiment.logdir is not None:
dump_environment(config, experiment.logdir, args.configs)
if self.distr_info:
info = yaml_load(self.task.additional_info)
info['resume'] = {
'master_computer': self.distr_info['master_computer'],
'master_task_id': self.task.id - self.distr_info['rank'],
'load_best': True
}
self.task.additional_info = yaml_dump(info)
self.task_provider.commit()
experiment.stages_config = {
k: v
for k, v in experiment.stages_config.items()
if k == experiment.stages[0]
}
runner.run_experiment(experiment, check=args.check)
if self.master and self.trace:
traced = trace_model_from_checkpoint(self.experiment.logdir, self)
torch.jit.save(traced, self.trace)
return {'stage': experiment.stages[-1], 'stages': stages}
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] = safitty.load(config_path)
runner_params = config.pop("runner_params", {}) or {}
# 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")
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_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