Tutorial 5: Customize Pruning algorithms¶
Here we show how to develop new Pruning algorithms with an example of AutoSlim.
Register a new algorithm
Create a new file
mmrazor/models/algorithms/autoslim.py
, classAutoSlim
inherits from classBaseAlgorithm
from mmrazor.models.builder import ALGORITHMS from .base import BaseAlgorithm @ALGORITHMS.register_module() class AutoSlim(BaseAlgorithm): def __init__(self, num_sample_training=4, input_shape=(3, 224, 224), bn_training_mode=False, **kwargs): super(AutoSlim, self).__init__(**kwargs) pass def train_step(self, data, optimizer): pass
Develop new algorithm components (optional)
AutoSlim can directly use class
RatioPruner
as core functions provider. If pruners provided in MMRazor don’t meet your needs, you can develop new algorithm components for your algorithm likeRatioPruner
. We will takeRatioPruner
as an example to introduce how to develop a new algorithm component:a. Create a new file
mmrazor/models/pruners/ratio_pruning.py
, classRatioPruner
can inherits from StructurePrunerb. Finish the functions you need, eg:
sample_subnet
,set_subnet
and so onfrom mmrazor.models.builder import PRUNERS, build_mutable from .structure_pruning import StructurePruner @PRUNERS.register_module() class RatioPruner(StructurePruner): def __init__(self, **kwargs): super().__init__(**kwargs) def sample_subnet(self): pass def set_subnet(self, subnet_dict): pass .... # supernet is a kind of architecture in `mmrazor/models/architectures/` def prepare_from_supernet(self, supernet): pass
c. Import the module in
mmrazor/models/pruners/__init__.py
from .ratio_pruning import RatioPruner __all__ = [..., 'RatioPruner']
Rewrite its train_step
Develop key logic of your algorithm in function
train_step
from mmrazor.models.builder import ALGORITHMS from .base import BaseAlgorithm @ALGORITHMS.register_module() class AutoSlim(BaseAlgorithm): def __init__(self, num_sample_training=4, input_shape=(3, 224, 224), bn_training_mode=False, **kwargs): super(AutoSlim, self).__init__(**kwargs) pass def train_step(self, data, optimizer): optimizer.zero_grad() losses = dict() if not self.retraining: ... else: ... optimizer.step() loss, log_vars = self._parse_losses(losses) outputs = dict( loss=loss, log_vars=log_vars, num_samples=len(data['img'].data)) return outputs
Add your custom functions (optional)
After finishing your key logic in function train_step, if you also need other custom functions, you can add them in class
AutoSlim
.Import the class
You can either add the following line to
mmrazor/models/algorithms/__init__.py
from .autoslim import AutoSlim __all__ = [..., 'AutoSlim']
Or alternatively add
custom_imports = dict( imports=['mmrazor.models.algorithms.spos'], allow_failed_imports=False)
to the config file to avoid modifying the original code.
Use the algorithm in your config file
algorithm = dict( type='AutoSlim', architecture=..., pruner=dict(type='RatioPruner', ...), retraining=...)