Announcing MMRazor: OpenMMLab Model Compression Toolbox and Benchmark.
MMRazor is an open source model compression toolbox based on PyTorch, which is the first toolbox that provides a framework for unified implementation and evaluation of model compression algorithms. It is a part of the OpenMMLab project.
Model compression has been an actively pursued area of research over the last few years with the goal of deploying state-of-the-art deep networks in low-power and resource-limited devices without significant drop in accuracy.
Parameter pruning, knowledge distillation, and neural architecture search are some suggested methods for compressing the size of deep networks. A model compression algorithm will use one or more of them，and can usually be applied to multiple models or tasks.
Yet, there has never been a codebase that can simultaneously support different methods and tasks, which significantly limits the research and application of model compression algorithms. Especially, the existing implementations of model compression algorithms are always coupled with multiple models, forcing researchers or developers to make relevant modification when applying an algorithm to different models, which is extremely user-unfriendly.
Here, we’re excited to announce our new project — MMRazor, which strives to overcome these challenges! Thanks to the generic framework from OpenMMLab, MMRazor is able to unify the implementation and evaluation of model compression algorithms. In addition, it has many more features, including but not limited to:
- All in one
MMRazor includes Neural Architecture Search（NAS), Knowledge Distillation（KD), Pruning, and Quantization（coming soon).And it supports the combination of different types of algorithms.
- General CV Model Compression ToolBox
Thanks to OpenMMLab, the algorithms in MMRazor can be quickly applied to different cv tasks, making the development of model compression algorithms once and for all.
- Decouple model and compression algorithm
MMRazor has a variety of built-in automation mechanisms, allowing developers to implement model compression algorithms without modifying the raw model code, such as:
OP can be modified code-free.
Feature maps in the middle layers can be obtained code-free.
Automatically obtain and analyze the connection between nn.Module.
And More …
- Flexible and Modular Design
We decompose the model compression algorithms into different components, making it much easier and more flexible to build a new algorithm by combining different components.
Below is an overview of MMRazor’s design and implementation
In terms of overall design, MMRazor mainly includes Component and Algorithm.
Component can be divided into basic component and algorithm component. Basic component consists of searcher, OP, Mutables and other modules in the figure, which provide basic function support for algorithm component. Algorithm component consists of Mutator, Pruner, Distiller and other modules in the figure. They provide core functionality for implementing various lightweight algorithms. The combination of Algorithm and Application can realize the purpose of slimming various task models.
In terms of implementation, MMRazor’s algorithm mainly contains two parts, namely architecture and algorithm components.
Architecture is similar to a model wrapper and can be easily combined with other OpenMMLab repos. Architecture plays different roles in different tasks, such as in one-Shot NAS, it is supernet; In distillation, it is student model; In pruning, it is a large model to be pruned; In quantization, it is the floating point model to be quantized.
Algorithm components can be flexibly called by each lightweight algorithm. Thanks to the componentization of the algorithm, it can be used alone or in combination in the specific lightweight algorithm. It only needs to rewrite train_step to realize the specific call logic of algorithm components.
The overall style of MMRazor is same as OpenMMLab’s . Both of them can flexibly configure the experiment through Config. Config mainly includes two parts, one is specific to MMRazor algorithm, and the other part can reuse the experimental configuration of other OpenMMLab repos. Including model definition, data processing, training schedule, etc.
Thanks to OpenMMLab’s powerful and highly flexible config mode and registry mechanism, MMRazor can perform ablation experiments by editing configuration files without changing the code.
If these are exactly what you need, don’t hesitate to check out our project on GitHub https://github.com/open-mmlab/mmrazor! Feel free to watch, star, and clone our projects. You are also welcome to help its growth by raising any issues and even making a PR!
