Call ofa_model() to get an OFA model. This
method finds all the nn.Conv2d / nn.ConvTranspose2d and nn.BatchNorm2d modules and replaces those modules with DynamicConv2d / DynamicConvTranspose2d and DynamicBatchNorm2d.
A list of pruning ratios is required to specify a maximum and minimum channel pruning ratio for each layer in the final OFA model.
An arbitrary pruning ratio is used in the output channel 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 in the list represent the subnetworks to be optimized. The pruning ratio is set by default to [0.5, 0.75, 1].
For a subnetwork sampled from the OFA model, the output channels of a
convolution layer are one of the numbers in the pruning ratio list multiplied by its
original number, for example, for a pruning ratio list of [0.5, 0.75, 1] and a
convolution layer nn.Conv2d(16, 32, 5), the output
channels of this layer in a sampled subnetwork is one of [0.5*32, 0.75*32, 1*32].
Because the first and last layers significantly impact network performance,
they are commonly excluded from pruning. By default, this method automatically
identifies the first and last convolutions, putting them into the list of excludes.
Setting auto_add_excludes equals False can cancel this
feature.
ofa_model = ofa_pruner.ofa_model([0.5, 0.75, 1], excludes = None, auto_add_excludes=True)