Algorithm¶
Introduction¶
What is algorithm in MMRazor¶
MMRazor is a model compression toolkit, which includes 4 mianstream technologies:
Neural Architecture Search (NAS)
Pruning
Knowledge Distillation (KD)
Quantization (come soon)
And in MMRazor, algorithm is a general item for these technologies. For example, in NAS,
SPOS is an algorithm, CWD is also an algorithm of knowledge distillation.
algorithm is the entrance of mmrazor/models . Its role in MMRazor is the same as both classifier in MMClassification and detector in MMDetection.
About base algorithm¶
In the directory of models/algorithms, all model compression algorithms are divided into 4 subdirectories: nas / pruning / distill / quantization. These algorithms must inherit from BaseAlgorithm, whose definition is as below.
from typing import Dict, List, Optional, Tuple, Union
from mmengine.model import BaseModel
from mmrazor.registry import MODELS
LossResults = Dict[str, torch.Tensor]
TensorResults = Union[Tuple[torch.Tensor], torch.Tensor]
PredictResults = List[BaseDataElement]
ForwardResults = Union[LossResults, TensorResults, PredictResults]
@MODELS.register_module()
class BaseAlgorithm(BaseModel):
def __init__(self,
architecture: Union[BaseModel, Dict],
data_preprocessor: Optional[Union[Dict, nn.Module]] = None,
init_cfg: Optional[Dict] = None):
......
super().__init__(data_preprocessor, init_cfg)
self.architecture = architecture
def forward(self,
batch_inputs: torch.Tensor,
data_samples: Optional[List[BaseDataElement]] = None,
mode: str = 'tensor') -> ForwardResults:
if mode == 'loss':
return self.loss(batch_inputs, data_samples)
elif mode == 'tensor':
return self._forward(batch_inputs, data_samples)
elif mode == 'predict':
return self._predict(batch_inputs, data_samples)
else:
raise RuntimeError(f'Invalid mode "{mode}". '
'Only supports loss, predict and tensor mode')
def loss(
self,
batch_inputs: torch.Tensor,
data_samples: Optional[List[BaseDataElement]] = None,
) -> LossResults:
"""Calculate losses from a batch of inputs and data samples."""
return self.architecture(batch_inputs, data_samples, mode='loss')
def _forward(
self,
batch_inputs: torch.Tensor,
data_samples: Optional[List[BaseDataElement]] = None,
) -> TensorResults:
"""Network forward process."""
return self.architecture(batch_inputs, data_samples, mode='tensor')
def _predict(
self,
batch_inputs: torch.Tensor,
data_samples: Optional[List[BaseDataElement]] = None,
) -> PredictResults:
"""Predict results from a batch of inputs and data samples with post-
processing."""
return self.architecture(batch_inputs, data_samples, mode='predict')
As you can see from above, BaseAlgorithm is inherited from BaseModel of MMEngine. BaseModel implements the basic functions of the algorithmic model, such as weights initialize,
batch inputs preprocess (see more information in BaseDataPreprocessor class of MMEngine), parse losses, and update model parameters. For more details of BaseModel , you can see docs for BaseModel.
BaseAlgorithm’s forward is just a wrapper of BaseModel’s forward. Sub-classes inherited from BaseAlgorithm only need to override the loss method, which implements the logic to calculate loss, thus various algorithms can be trained in the runner.
How to use existing algorithms in MMRazor¶
Configure your architecture that will be slimmed
Use the model config of other repos of OpenMMLab directly as below, which is an example of setting Faster-RCNN as our architecture.
_base_ = [
'mmdet::_base_/models/faster_rcnn_r50_fpn.py',
]
architecture = _base_.model
Use your customized model as below, which is an example of defining a VGG model as our architecture.
Note
How to customize architectures can refer to our tutorial: Customize Architectures.
default_scope='mmcls'
architecture = dict(
type='ImageClassifier',
backbone=dict(type='VGG', depth=11, num_classes=1000),
neck=None,
head=dict(
type='ClsHead',
loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
topk=(1, 5),
))
Apply the registered algorithm to your architecture.
Note
The arg name of algorithm in config is model rather than algorithm in order to get better supports of MMCV and MMEngine.
Maybe more args in model need to set according to the used algorithm.
model = dict(
type='BaseAlgorithm',
architecture=architecture)
Note
About the usage of Config, refer to config.md please.
Apply some custom hooks or loops to your algorithm. (optional)
Custom hooks
custom_hooks = [
dict(type='NaiveVisualizationHook', priority='LOWEST'),
]
Custom loops
_base_ = ['./spos_shufflenet_supernet_8xb128_in1k.py']
# To chose from ['train_cfg', 'val_cfg', 'test_cfg'] based on your loop type
train_cfg = dict(
_delete_=True,
type='mmrazor.EvolutionSearchLoop',
dataloader=_base_.val_dataloader,
evaluator=_base_.val_evaluator)
val_cfg = dict()
test_cfg = dict()
How to customize your algorithm¶
Common pipeline¶
Register a new algorithm
Create a new file mmrazor/models/algorithms/{subdirectory}/xxx.py
from mmrazor.models.algorithms import BaseAlgorithm
from mmrazor.registry import MODELS
@MODELS.register_module()
class XXX(BaseAlgorithm):
def __init__(self, architecture):
super().__init__(architecture)
pass
def loss(self, batch_inputs):
pass
Rewrite its
lossmethod.
from mmrazor.models.algorithms import BaseAlgorithm
from mmrazor.registry import MODELS
@MODELS.register_module()
class XXX(BaseAlgorithm):
def __init__(self, architecture):
super().__init__(architecture)
......
def loss(self, batch_inputs):
......
return LossResults
Add the remaining functions of the algorithm
Note
This step is special because of the diversity of algorithms. Some functions of the algorithm may also be implemented in other files.
from mmrazor.models.algorithms import BaseAlgorithm
from mmrazor.registry import MODELS
@MODELS.register_module()
class XXX(BaseAlgorithm):
def __init__(self, architecture):
super().__init__(architecture)
......
def loss(self, batch_inputs):
......
return LossResults
def aaa(self):
......
def bbb(self):
......
Import the class
You can add the following line to mmrazor/models/algorithms/{subdirectory}/__init__.py
from .xxx import XXX
__all__ = ['XXX']
In addition, import XXX in mmrazor/models/algorithms/__init__.py
Use the algorithm in your config file.
Please refer to the previous section about how to use existing algorithms in MMRazor
model = dict(
type='XXX',
architecture=architecture)
Pipelines for different algorithms¶
Please refer to our tutorials about how to customize different algorithms for more details as below.
NAS
Pruning
Distill