pytorch_nndct.OFAPruner - 2.0 English

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 re-calibrated 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 FLOPs-accuracy trade-offs 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.