Tutorial 5: Customize Pruning algorithms

Here we show how to develop new Pruning algorithms with an example of AutoSlim.

1. Register a new algorithm

   Create a new file mmrazor/models/algorithms/autoslim.py, class AutoSlim inherits from class BaseAlgorithm

   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
   

2. 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 like RatioPruner. We will take RatioPruner as an example to introduce how to develop a new algorithm component:

   a. Create a new file mmrazor/models/pruners/ratio_pruning.py, class RatioPruner can inherits from StructurePruner

   b. Finish the functions you need, eg: sample_subnet, set_subnet and so on

   from 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']
   

3. Rewrite its train_step

   Develop key logic of your algorithm in functiontrain_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
   

4. 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.

5. 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.

6. Use the algorithm in your config file

   algorithm = dict(
       type='AutoSlim',
       architecture=...,
       pruner=dict(type='RatioPruner', ...),
       retraining=...)