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)
 expand_ratio
 A list of prune ratio of each convolution layer (kernel > 1). For
each convolution layer in the OFA model, arbitrary pruning ratio can
be used in output channel.
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.

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 BatchNormalization layers.
 model
 The OFA model.

run_evolutionary_search(model, calibration_fn, calib_args, eval_fn, eval_args, min_flops, max_flops, flops_step=10, parent_popu_size=16, iteration=0 , utate_size=8 ,mutate_prob=0.2, crossover_size=84
Runs an evolutionary search to find subnetworks whose flops 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 for evaluation.
 calib_args
 The arguments for calibration_fn.
 eval_fn
 A function to evaluate the model.
 eval_args
 The arguments for eval_fn.
 min_flops
 The minimum flops of searched subnetworks.
 max_flops
 The maximum flops of searched subnetworks.
 flops_step
 The step of flops for searching. Divides the interval [min_flops, max flops] into segments by flops_step. For each segment, searches the best FLOPsaccuracy tradeoffs subnetwork.
 parent_popu_size
 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.
 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(self, 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.