Pruning

Runners

class compressors.pruning.runners.PruneRunner(num_sessions: int = 5, input_key: Union[str, List[str]] = 'features', output_key: Union[str, List[str]] = 'logits', *runner_args, **runner_kwargs)
handle_batch(batch: Mapping[str, Any])None

Inner method to handle specified data batch. Used to make a train/valid/infer stage during Experiment run.

Parameters

batch (Mapping[str, Any]) – dictionary with data batches from DataLoader.

predict_batch(batch: Mapping[str, Any], **kwargs)Mapping[str, Any]

Run model inference on specified data batch.

Parameters
  • batch – dictionary with data batches from DataLoader.

  • **kwargs – additional kwargs to pass to the model

Returns

model output dictionary

Return type

Mapping

Raises

NotImplementedError – if not implemented yet

property stages

Experiment’s stage names (array with one value).

Callbacks

class compressors.pruning.callbacks.LotteryTicketCallback(initial_state_dict: Dict[str, Any])
on_stage_start(runner: catalyst.core.runner.IRunner)None

Event handler.

Parameters

runner – experiment runner

class compressors.pruning.callbacks.PrepareForFinePruningCallback(train_loader_key: str = 'train', *logits_dataset_args, **logits_dataset_kwargs)
on_experiment_start(runner: catalyst.core.runner.IRunner)None

Event handler for experiment start.

Parameters

runner – IRunner instance.

Note

This event work only on IRunner.

Utils