puzzletron

This module provides the main compression function for a model using MIP-based NAS search algorithm.

Functions

puzzletron

Compress a model using the MIP-based NAS search algorithm from Puzzletron.
puzzletron(hydra_config_dir, hydra_config, puzzle_dir, dataset_path)

Compress a model using the MIP-based NAS search algorithm from Puzzletron.

Parameters:

  • hydra_config_dir (str) – path to a hydra_config_dir that defines the search space

  • hydra_config (str) – the corresponding hydra config file

  • puzzle_dir (str) – directory with a puzzletron model to compress

  • dataset_path (str) – dataset used for scoring and distillation

Returns:

Hydra config object after compressing the model. The same hydra configuration object is used across all compression steps. TODO: Investigate if this config object is immutable across steps and clarify

Return type:

DictConfig