This API has the following methods:
torch.nn.Moduleobject to prune.
- 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 the real data.
ofa_model(expand_ratio, channel_divisble=8, excludes=None, auto_add_excludes=True, save_search_space=False)
- A list of prune ratio of each convolution layer. For each convolution
layer in the OFA model, arbitrary pruning ratio can be used in
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. By default, the pruning ratio is set to [0.5, 0.75, 0.1].
- A channel number that is divisible by a given divisor.
- A list of modules to be excluded from pruning.
- Bool. If True, this method automatically identifies the first convolution and the last convolution, then put them into the list of excludes. If False, not doing the above. Defaults to True.
- Bool. If True, save the search space of the model as a file of 'searchspace.config'. Customers can check the search space for each layer. Defaults to False.
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.
- The OFA model.
- One of ['random', 'max', 'min'].
Resets the running stats of the Batch Normalization layers.
- The OFA model.
run_evolutionary_search(model, calibration_fn, calib_args, eval_fn, eval_args, evaluation_metric, min_or_max_metric, min_flops, max_flops, flops_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 flops are in a given range.
- The OFA model.
- 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 for evaluation.
- The arguments for calibration_fn.
- A function to evaluate the model.
- The arguments for eval_fn.
- A string of evaluation_metric to record the result.
- 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, when the evaluation metric is mean squared error (mse) or mean absolute error (mae), record the minimum value.
- The minimum flops of searched subnetworks.
- The maximum flops of searched subnetworks.
- The step of flops for searching. Divides the interval [min_flops, max flops] into segments by flops_step. For each segment, searches the best FLOPs-accuracy trade-offs subnetwork.
- The number of initial parent population for sampling the given number of random subnetworks whose flops are in the given range. The larger this number is, the longer the search will take and the more likely it is that the best results will be obtained.
- The number of iterations for searching or the number of cycles of the whole algorithm.
- The size of mutation. Each value of the subnetwork setting is replaced with another value of the candidate list with probability of mutate_prob.
- The probability of mutation.
- The size of crossover. Sampling two subnetwork settings and swapping any value in the two subnetwork settings randomly.
Saving dynamic/static subnetwork settings with JSON.
- The configurations for the dynamic subnetwork setting.
- Filepath to save the subnetwork settings.
- Filepath to load the subnetwork settings.