This API has the following methods:
-
__init__(model, inputs)- model
- A
torch.nn.Moduleobject to prune. - inputs
- A single or a list of torch. Tensor used as inputs for model inference. It does not need to be real data. It can be a randomly generated tensor of the same shape and data type as real data.
-
ofa_model(expand_ratio, channel_divisble=8, excludes=None, auto_add_excludes=True, save_search_space=False)- expand_ratio
- A list of prune ratio of each convolution layer. An arbitrary pruning
ratio can be used for the output channels for each convolution layer
in the OFA model.
The maximum and minimum values in this list represent the maximum and minimum compression rates of the model. Other values represent subnetworks to be optimized. The pruning ratio is set by default to [0.5, 0.75, 0.1].
- channel_divisible
- A channel number that is divisible by a given divisor.
- excludes
- A list of modules to be excluded from pruning.
- auto_add_excludes
- Bool. If True, this method automatically identifies the first and the last convolution and puts them into the list of excludes. If False, skips creating the list of excludes. Defaults to True.
- save_search_space
- Bool. If True, save the search space of the model as a file of 'searchspace.config'. You can check the search space for each layer. Defaults to False.
-
sample_subnet(model, mode)Returns a subnetwork and its configuration for a given mode. The subnetwork can do a forward/backward process using a part of the weights from the OFA model and its settings.
- model
- The OFA model.
- mode
- One of ['random', 'max', 'min'].
-
reset_bn_running_stats_for_calibration(model)Resets the running stats of the Batch Normalization layers.
- model
- The OFA model.
-
run_evolutionary_search(model, calibration_fn, calib_args, eval_fn, eval_args, evaluation_metric, min_or_max_metric, min_macs, max_macs, macs_step=10, parent_popu_size=16, iteration=10, mutate_size=8, mutate_prob=0.2, crossover_size=4)Runs an evolutionary search to find the best subnetwork whose MACs are in a given range.
- model
- The OFA model.
- calibration_fn
- A BatchNormalization calibration function. All subnetworks share weights in an OFA model, but batch normalization statistics (mean and variance) are not stored when training the OFA model. After the training is completed, the batch normalization statistics must be re-calibrated using the training data for each sampled subnetwork used for evaluation.
- calib_args
- The arguments for calibration_fn.
- eval_fn
- A function to evaluate the model.
- eval_args
- The arguments for eval_fn.
- evaluation_metric
- A string of evaluation_metric to record the result.
- min_or_max_metric
- One of ['max', 'min']. The maximum or minimum value of the evaluation metric to be recorded in the evolutionary search. For example, when the evaluation metric has an accuracy of top1, record the maximum value of each iteration in the evolutionary search. However, record the minimum value when the evaluation metric is a mean squared error (mse) or mean absolute error (mae).
- min_macs
- The minimum MACs of searched subnetworks.
- max_macs
- The maximum MACs of searched subnetworks.
- macs_step
- The step of MACs for searching. Divide the interval [min_macs, max macs] into segments by macs_step. For each segment, search the best macs-accuracy trade-offs subnetwork.
- parent_popu_size
- The number of initial parent population for sampling the given number of random subnetworks whose MACs are in the given range. The larger this number is, the longer the search takes, and the more likely the best results are obtained.
- iteration
- The number of iterations for searching or the number of cycles of the whole algorithm.
- mutate_size
- The size of mutation. Each value of the subnetwork setting is replaced with another value of the candidate list with a probability of mutate_prob.
- mutate_prob
- The probability of mutation.
- crossover_size
- The size of the crossover. Sampling two subnetwork settings and swapping any value in the two subnetwork settings randomly.
-
save_subnet_config(setting_config, file_name)Saving dynamic/static subnetwork settings with JSON.
- setting_config
- The configurations for the dynamic subnetwork setting.
- file_name
- Filepath to save the subnetwork settings.
-
load_subnet_config(file_name)- file_name
- Filepath to load the subnetwork settings.