This document introduces the pruning framework in mmrazor. Our pruning framework can help you prune a model automatically, making it easy to extend new algorithms.
The pruning framework consists of five modules: Algorithm, ChanelMutator, MutableChannelUnit, MutableChannel, and DynamicOp. Their main features are detailed below:
A ChanelMutator controls the pruning structure of a model. In other words, ChanelMutator decides how many channels each layer prunes. Usually, given a pruning target, such as a flops, latency, or pruning ratio target, the ChannelUnitMutator will output a pruning structure for the model. The pruning structure is variable. The default definition is the remaining channel ratio, and it's also easy to extend to the number of channels or channel buckets.
As some layers' channels are related, the related layers share one pruning decision. We put these associated layers into a MutableChannelUnit. Therefore, the ChanelMutator directly decides the pruning ratio of each MutableChannelUnit.
```python
from mmrazor.models.mutators import BaseChannelMutator
from mmengine.model import BaseModel
import torch.nn as nn
class Model(BaseModel):
def __init__(self):
super().__init__()
self.feature = nn.Sequential(
nn.Conv2d(3, 8, 3, 2, 1),
nn.Conv2d(8, 16, 3, 2, 1)
)
self.pool = nn.AdaptiveAvgPool2d(1)
self.head = nn.Linear(16, 1000)
def forward(self, x):
x_ = self.pool(self.feature(x))
return self.head(x_.flatten(1))
model = Model()
mutator = BaseChannelMutator()
mutator.prepare_from_supernet(model)
print(mutator.sample_choices())
# {
# 'feature.0_(0, 8)_out_1_in_1': 0.5,
# 'feature.1_(0, 16)_out_1_in_1': 0.5625
# }
```
Please refer to [ChannelMutator](../../../mmrazor/models/mutables/mutable_channel/units/mutable_channel_unit.ipynb) for more details.
**MutableChannel**: Each MutableChannel manages a channel mask for a model. They help DynamicOps to deal with mutable numbers of channels. Please refer to [MutableChannel](../../../mmrazor/models/mutables/mutable_channel/MutableChannel.md) for more details.
**DynamicOp**: DynamicOps inherit from basic torch modules, like nn.Conv2d or nn.Linear. They can forward with mutable numbers of channels and export pruned torch modules.
Compared with basic torch modules, each DynamicOp has two MutableChannel modules, which control the input and output channels.
## More Documents about Pruning
Please refer to the following documents for more details.
- Development tutorials
- [How to prune your model](../advanced_guides/tutorials/how_to_prune_your_model.md)
- [How to use config tool of pruning](../advanced_guides/tutorials/how_to_use_config_tool_of_pruning.md)
