This API has the following methods:

__init__(model, inputs)
 model
 A
torch.nn.Module
object 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 the 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. For each convolution
layer in the OFA model, an arbitrary pruning ratio can be used for
the output channels.
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].
 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 convolution and the last convolution, then put them into the list of excludes. If False, not doing the above. 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 recalibrated 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, when the evaluation metric is mean squared error (mse) or mean absolute error (mae), record the minimum value.
 min_macs
 The minimum macs of searched subnetworks.
 max_macs
 The maximum macs of searched subnetworks.
 macs_step
 The step of macs for searching. Divides the interval [min_macs, max macs] into segments by macs_step. For each segment, searches the best macsaccuracy tradeoffs 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 will take and more likely that the best results will be 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 probability of mutate_prob.
 mutate_prob
 The probability of mutation.
 crossover_size
 The size of 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.