From c650b3eee56121258e0a5d8d16105d17e22234f3 Mon Sep 17 00:00:00 2001 From: PJDong <1115957667@qq.com> Date: Tue, 30 Aug 2022 19:48:18 +0800 Subject: [PATCH] [Doc] Add Tutorial of KD,Pruning,NAS and Installation. (#255) * [doc] add doc of installation, kd, nas, pruning * [doc] add user guide part * fix name of mmcv-full to mmcv * [doc] fix mmcv based on zaida 's suggestion * [doc] fix index error based on shiguang's comments --- docs/en/get_started/installation.md | 110 +++++++++++++ .../2_train_different_types_algorithms.md | 106 ++++++++++++ .../3_train_with_different_devices.md | 92 +++++++++++ docs/en/user_guides/4_test_a_model.md | 78 +++++++++ docs/en/user_guides/kd.md | 123 ++++++++++++++ docs/en/user_guides/nas.md | 135 ++++++++++++++++ docs/en/user_guides/pruning.md | 151 ++++++++++++++++++ 7 files changed, 795 insertions(+) diff --git a/docs/en/get_started/installation.md b/docs/en/get_started/installation.md index 25267fe2b..19a483593 100644 --- a/docs/en/get_started/installation.md +++ b/docs/en/get_started/installation.md @@ -1 +1,111 @@ # Installation + +## Prepare Environment + +Create a conda virtual environment and activate it. + +```Python +conda create -n openmmlab python=3.7 -y +conda activate openmmlab +``` + +Install PyTorch and torchvision following the [official instructions](https://pytorch.org/). + +Note: Make sure that your compilation CUDA version and runtime CUDA version match. You can check the supported CUDA version for precompiled packages on the [PyTorch website](https://pytorch.org/). + +`E.g.1` If you have CUDA 10.2 installed under `/usr/local/cuda` and would like to install PyTorch 1.10, you need to install the prebuilt PyTorch with CUDA 10.2. + +```Python +conda install pytorch torchvision torchaudio cudatoolkit=10.2 -c pytorch +``` + +`E.g.2` If you have CUDA 9.2 installed under `/usr/local/cuda` and would like to install PyTorch 1.5.1, you need to install the prebuilt PyTorch with CUDA 9.2. + +```Python +conda install pytorch==1.5.1 torchvision==0.6.1 cudatoolkit=9.2 -c pytorch +``` + +- If you build PyTorch from source instead of installing the prebuilt package, you can use more CUDA versions such as 9.0. + +## Customize Installation + +It is recommended to install MMRazor with [MIM](https://github.com/open-mmlab/mim), which automatically handles the dependencies of OpenMMLab projects, including mmcv and other python packages. + +```Python +pip install openmim +mim install git+https://github.com/open-mmlab/mmrazor.git@1.0.0rc0 +``` + +Or you can still install MMRazor manually + +1. Install mmcv. + +```Python +pip install 'mmcv>=2.0.0rc1' -f https://download.openmmlab.com/mmcv/dist/{cu_version}/{torch_version}/index.html +``` + +Please replace `{cu_version}` and `{torch_version}` in the url to your desired one. For example, to install the latest `mmcv` with `CUDA 10.2` and `PyTorch 1.10.0`, use the following command: + +```Python +pip install 'mmcv>=2.0.0rc1' -f https://download.openmmlab.com/mmcv/dist/cu102/torch1.10.0/index.html +``` + +See [here](https://github.com/open-mmlab/mmcv#installation) for different versions of MMCV compatible to different PyTorch and CUDA versions. + +Optionally, you can compile mmcv from source. + +``` +MMCV_WITH_OPS=0 pip install -e . -v +# install mmcv-lite, do not compile operators +MMCV_WITH_OPS=1 pip install -e . -v +# install mmcv (originally called mmcv-full), compile operators +pip install -e . -v +# install mmcv with compiled operators, +``` + +2. Install MMEngine. + +Compile MMEngine from source. + +```Python +git clone https://github.com/open-mmlab/mmengine.git +cd mmengine +pip install -v -e . +``` + +3. Install MMRazor. + +If you would like to install MMRazor in `dev` mode, run following: + +```Python +git clone https://github.com/open-mmlab/mmrazor.git +cd mmrazor +git fetch origin +git checkout -b 1.0.0rc0 origin/1.0.0rc0 +# The new version is released in branch ``1.0.0rc0`` +pip install -v -e . +# "-v" means verbose, or more output +# "-e" means installing a project in editable mode, +# thus any local modifications made to the code will take effect without reinstallation. +``` + +**Note:** + +- When MMRazor is installed on `dev` mode, any local modifications made to the code will take effect without the need to reinstall it. + +## A from-scratch Setup Script + +```Python +conda create -n openmmlab python=3.7 -y +conda activate openmmlab + +conda install pytorch torchvision cudatoolkit=10.2 -c pytorch +# install the latest mmcv +pip install 'mmcv>=2.0.0rc1' -f https://download.openmmlab.com/mmcv/dist/cu102/torch1.10.0/index.html +# install mmrazor +git clone https://github.com/open-mmlab/mmrazor.git +cd mmrazor +git fetch origin +git checkout -b 1.0.0rc0 origin/1.0.0rc0 +pip install -v -e . +``` diff --git a/docs/en/user_guides/2_train_different_types_algorithms.md b/docs/en/user_guides/2_train_different_types_algorithms.md index 61385b4b0..01f295581 100644 --- a/docs/en/user_guides/2_train_different_types_algorithms.md +++ b/docs/en/user_guides/2_train_different_types_algorithms.md @@ -1 +1,107 @@ # Train different types algorithms + +**Before running our algorithms, you may need to prepare the datasets according to the instructions in the corresponding document.** + +**Note**: + +- With the help of mmengine, mmrazor unified entered interfaces for various tasks, thus our algorithms will adapt all OpenMMLab upstream repos in theory. + +- We dynamically pass arguments `cfg-options` (e.g., `mutable_cfg` in nas algorithm or `channel_cfg` in pruning algorithm) to **avoid the need for a config for each subnet or checkpoint**. If you want to specify different subnets for retraining or testing, you just need to change this argument. + +### NAS + +Here we take SPOS(Single Path One Shot) as an example. There are three steps to start neural network search(NAS), including **supernet pre-training**, **search for subnet on the trained supernet** and **subnet retraining**. + +#### Supernet Pre-training + +```Python +python tools/train.py ${CONFIG_FILE} [optional arguments] +``` + +The usage of optional arguments are the same as corresponding tasks like mmclassification, mmdetection and mmsegmentation. + +For example, + +```Python +python ./tools/train.py \ + configs/nas/mmcls/spos/spos_shufflenet_supernet_8xb128_in1k.py + --work-dir $WORK_DIR +``` + +#### Search for Subnet on The Trained Supernet + +```Python +python tools/train.py ${CONFIG_FILE} ${CHECKPOINT_PATH} [optional arguments] +``` + +For example, + +```Python +python ./tools/train.py \ + configs/nas/mmcls/spos/spos_shufflenet_search_8xb128_in1k.py \ + $STEP1_CKPT \ + --work-dir $WORK_DIR +``` + +#### Subnet Retraining + +```Python +python tools/train.py ${CONFIG_FILE} \ + --cfg-options algorithm.fix_subnet=${MUTABLE_CFG_PATH} [optional arguments] +``` + +- `MUTABLE_CFG_PATH`: Path of `fix_subnet`. `fix_subnet` represents **config for mutable of the subnet searched out**, used to specify different subnets for retraining. An example for `fix_subnet` can be found [here](https://github.com/open-mmlab/mmrazor/blob/master/configs/nas/spos/SPOS_SHUFFLENETV2_330M_IN1k_PAPER.yaml), and the usage can be found [here](https://github.com/open-mmlab/mmrazor/blob/master/configs/nas/spos/README.md#subnet-retraining-on-imagenet). + +For example, + +```Python +python ./tools/train.py \ + configs/nas/mmcls/spos/spos_shufflenet_subnet_8xb128_in1k.py \ + --work-dir $WORK_DIR \ + --cfg-options algorithm.fix_subnet=$YAML_FILE_BY_STEP2 +``` + +We note that instead of using `--cfg-options`, you can also directly modify ``` configs/nas/mmcls/spos/``spos_shufflenet_subnet_8xb128_in1k``.py ``` like this: + +```Python +fix_subnet = 'configs/nas/mmcls/spos/SPOS_SHUFFLENETV2_330M_IN1k_PAPER.yaml' +model = dict(fix_subnet=fix_subnet) +``` + +### Pruning + +Pruning has three steps, including **supernet pre-training**, **search for subnet on the trained supernet** and **subnet retraining**. The commands of the first two steps are similar to NAS, except that we need to use `CONFIG_FILE` of Pruning here. The commands of the **subnet retraining** are as follows. + +#### Subnet Retraining + +```Python +python tools/train.py ${CONFIG_FILE} --cfg-options model._channel_cfg_paths=${CHANNEL_CFG_PATH} [optional arguments] +``` + +Different from NAS, the argument that needs to be specified here is `channel_cfg_paths` . + +- `CHANNEL_CFG_PATH`: Path of `_channel_cfg_path`. `channel_cfg` represents **config for channel of the subnet searched out**, used to specify different subnets for testing. + +For example, the default `_channel_cfg_paths` is set in the config below. + +```Python +python ./tools/train.py \ + configs/pruning/mmcls/autoslim/autoslim_mbv2_1.5x_subnet_8xb256_in1k_flops-530M \ + --work-dir your_work_dir +``` + +### Distillation + +There is only one step to start knowledge distillation. + +```Python +python tools/train.py ${CONFIG_FILE} [optional arguments] +``` + +For example, + +```Python +python ./tools/train.py \ + configs/distill/mmcls/kd/kd_logits_r34_r18_8xb32_in1k.py \ + --work-dir your_work_dir +``` diff --git a/docs/en/user_guides/3_train_with_different_devices.md b/docs/en/user_guides/3_train_with_different_devices.md index 60d52fe56..aa24654e8 100644 --- a/docs/en/user_guides/3_train_with_different_devices.md +++ b/docs/en/user_guides/3_train_with_different_devices.md @@ -1 +1,93 @@ # Train with different devices + +**Note**: The default learning rate in config files is for 8 GPUs. If using different number GPUs, the total batch size will change in proportion, you have to scale the learning rate following `new_lr = old_lr * new_ngpus / old_ngpus`. We recommend to use `tools/dist_train.sh` even with 1 gpu, since some methods do not support non-distributed training. + +### Training with CPU + +```Python +export CUDA_VISIBLE_DEVICES=-1 +python tools/train.py ${CONFIG_FILE} +``` + +**Note**: We do not recommend users to use CPU for training because it is too slow and some algorithms are using `SyncBN` which is based on distributed training. We support this feature to allow users to debug on machines without GPU for convenience. + +### Train with single/multiple GPUs + +```Python +sh tools/dist_train.sh ${CONFIG_FILE} ${GPUS} --work_dir ${YOUR_WORK_DIR} [optional arguments] +``` + +**Note**: During training, checkpoints and logs are saved in the same folder structure as the config file under `work_dirs/`. Custom work directory is not recommended since evaluation scripts infer work directories from the config file name. If you want to save your weights somewhere else, please use symlink, for example: + +```Python +ln -s ${YOUR_WORK_DIRS} ${MMRAZOR}/work_dirs +``` + +Alternatively, if you run MMRazor on a cluster managed with [slurm](https://slurm.schedmd.com/): + +```Python +GPUS_PER_NODE=${GPUS_PER_NODE} GPUS=${GPUS} SRUN_ARGS=${SRUN_ARGS} sh tools/xxx/slurm_train.sh ${PARTITION} ${JOB_NAME} ${CONFIG_FILE} ${YOUR_WORK_DIR} [optional arguments] +``` + +### Train with multiple machines + +If you launch with multiple machines simply connected with ethernet, you can simply run the following commands: + +On the first machine: + +```Python +NNODES=2 NODE_RANK=0 PORT=$MASTER_PORT MASTER_ADDR=$MASTER_ADDR sh tools/dist_train.sh $CONFIG $GPUS +``` + +On the second machine: + +```Python +NNODES=2 NODE_RANK=1 PORT=$MASTER_PORT MASTER_ADDR=$MASTER_ADDR sh tools/dist_train.sh $CONFIG $GPUS +``` + +Usually it is slow if you do not have high speed networking like InfiniBand. + +If you launch with slurm, the command is the same as that on single machine described above, but you need refer to [slurm_train.sh](https://github.com/open-mmlab/mmselfsup/blob/master/tools/slurm_train.sh) to set appropriate parameters and environment variables. + +### Launch multiple jobs on a single machine + +If you launch multiple jobs on a single machine, e.g., 2 jobs of 4-GPU training on a machine with 8 GPUs, you need to specify different ports (29500 by default) for each job to avoid communication conflict. + +If you use `dist_train.sh` to launch training jobs: + +```Python +CUDA_VISIBLE_DEVICES=0,1,2,3 PORT=29500 sh tools/xxx/dist_train.sh ${CONFIG_FILE} 4 --work_dir tmp_work_dir_1 +CUDA_VISIBLE_DEVICES=4,5,6,7 PORT=29501 sh tools/xxx/dist_train.sh ${CONFIG_FILE} 4 --work_dir tmp_work_dir_2 +``` + +If you use launch training jobs with slurm, you have two options to set different communication ports: + +Option 1: + +In `config1.py`: + +```Python +dist_params = dict(backend='nccl', port=29500) +``` + +In `config2.py`: + +```Python +dist_params = dict(backend='nccl', port=29501) +``` + +Then you can launch two jobs with config1.py and config2.py. + +```Python +CUDA_VISIBLE_DEVICES=0,1,2,3 GPUS=4 sh tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config1.py tmp_work_dir_1 +CUDA_VISIBLE_DEVICES=4,5,6,7 GPUS=4 sh tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config2.py tmp_work_dir_2 +``` + +Option 2: + +You can set different communication ports without the need to modify the configuration file, but have to set the `cfg-options` to overwrite the default port in configuration file. + +```Python +CUDA_VISIBLE_DEVICES=0,1,2,3 GPUS=4 sh tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config1.py tmp_work_dir_1 --cfg-options dist_params.port=29500 +CUDA_VISIBLE_DEVICES=4,5,6,7 GPUS=4 sh tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config2.py tmp_work_dir_2 --cfg-options dist_params.port=29501 +``` diff --git a/docs/en/user_guides/4_test_a_model.md b/docs/en/user_guides/4_test_a_model.md index 7a61fd7fd..152b5dc41 100644 --- a/docs/en/user_guides/4_test_a_model.md +++ b/docs/en/user_guides/4_test_a_model.md @@ -1 +1,79 @@ # Test a model + +### NAS + +To test nas method, you can use the following command. + +```Python +python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_PATH} --cfg-options algorithm.fix_subnet=${FIX_SUBNET_PATH} [optional arguments] +``` + +- `FIX_SUBNET_PATH`: Path of `fix_subnet`. `fix_subnet` represents **config for mutable of the subnet searched out**, used to specify different subnets for testing. An example for `fix_subnet` can be found [here](https://github.com/open-mmlab/mmrazor/blob/master/configs/nas/spos/SPOS_SHUFFLENETV2_330M_IN1k_PAPER.yaml). + +The usage of optional arguments are the same as corresponding tasks like mmclassification, mmdetection and mmsegmentation. + +For example, + +```Python +python tools/test.py \ + configs/nas/mmcls/spos/spos_subnet_shufflenetv2_8xb128_in1k.py \ + your_subnet_checkpoint_path \ + --cfg-options algorithm.fix_subnet=configs/nas/mmcls/spos/SPOS_SHUFFLENETV2_330M_IN1k_PAPER.yaml +``` + +### Pruning + +#### Split Checkpoint(Optional) + +If you train a slimmable model during retraining, checkpoints of different subnets are actually fused in only one checkpoint. You can split this checkpoint to multiple independent checkpoints by using the following command + +```Python +python tools/model_converters/split_checkpoint.py ${CONFIG_FILE} ${CHECKPOINT_PATH} --channel-cfgs ${CHANNEL_CFG_PATH} [optional arguments] +``` + +- `CHANNEL_CFG_PATH`: A list of paths of `channel_cfg`. For example, when you retrain a slimmable model, your command will be like `--cfg-options algorithm.channel_cfg=cfg1,cfg2,cfg3`. And your command here should be `--channel-cfgs cfg1 cfg2 cfg3`. The order of them should be the same. + +For example, + +```Python +python tools/model_converters/split_checkpoint.py \ + configs/pruning/autoslim/autoslim_mbv2_subnet_8xb256_in1k.py \ + your_retraining_checkpoint_path \ + --channel-cfgs configs/pruning/autoslim/AUTOSLIM_MBV2_530M_OFFICIAL.yaml configs/pruning/autoslim/AUTOSLIM_MBV2_320M_OFFICIAL.yaml configs/pruning/autoslim/AUTOSLIM_MBV2_220M_OFFICIAL.yaml +``` + +#### Test + +To test pruning method, you can use following command + +```Python +python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_PATH} --cfg-options model._channel_cfg_paths=${CHANNEL_CFG_PATH} [optional arguments] +``` + +- `task`: one of `mmcls`、`mmdet` and `mmseg` + +- `CHANNEL_CFG_PATH`: Path of `channel_cfg`. `channel_cfg` represents **config for channel of the subnet searched out**, used to specify different subnets for testing. An example for `channel_cfg` can be found [here](https://github.com/open-mmlab/mmrazor/blob/master/configs/pruning/autoslim/AUTOSLIM_MBV2_220M_OFFICIAL.yaml), and the usage can be found [here](https://github.com/open-mmlab/mmrazor/blob/master/configs/pruning/autoslim/README.md#test-a-subnet). + +For example, + +```Python +python ./tools/test.py \ + configs/pruning/mmcls/autoslim/autoslim_mbv2__1.5x_subnet_8xb256_in1k-530M.py \ + your_splitted_checkpoint_path --metrics accuracy +``` + +### Distillation + +To test the distillation method, you can use the following command + +```Python +python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_PATH} [optional arguments] +``` + +For example, + +```Python +python ./tools/test.py \ + configs/distill/mmseg/cwd/cwd_logits_pspnet_r101_d8_pspnet_r18_d8_512x1024_cityscapes_80k.py \ + your_splitted_checkpoint_path --show +``` diff --git a/docs/en/user_guides/kd.md b/docs/en/user_guides/kd.md index be95c5f18..072732fbe 100644 --- a/docs/en/user_guides/kd.md +++ b/docs/en/user_guides/kd.md @@ -1 +1,124 @@ # Knowledge Distillation + +Here we show how to develop new KD algorithms with an example of `SingleTeacherDistill`. + +1. Register a new algorithm + +Create a new file `mmrazor/models/algorithms/distill/configurable/single_teacher_distill.py`, class `SingleTeacherDistill` inherits from class `BaseAlgorithm` + +```Python +from mmrazor.registry import MODELS +from ..base import BaseAlgorithm + +@ALGORITHMS.register_module() +class SingleTeacherDistill(BaseAlgorithm): + def __init__(self, use_gt, **kwargs): + super(Distillation, self).__init__(**kwargs) + pass + + def train_step(self, data, optimizer): + pass +``` + +2. Develop connectors (Optional) . + +Take ConvModuleConnector as an example. + +```python +from mmrazor.registry import MODELS +from .base_connector import BaseConnector + +@MODELS.register_module() +class ConvModuleConncetor(BaseConnector): + def __init__(self, in_channel, out_channel, kernel_size = 1, stride = 1): + ... + + def forward_train(self, feature): + ... +``` + +3. Develop distiller. + +Take `ConfigurableDistiller` as an example. + +```python +from .base_distiller import BaseDistiller +from mmrazor.registry import MODELS + + +@MODELS.register_module() +class ConfigurableDistiller(BaseDistiller): + def __init__(self, + student_recorders = None, + teacher_recorders = None, + distill_deliveries = None, + connectors = None, + distill_losses = None, + loss_forward_mappings = None): + ... + + def build_connectors(self, connectors): + ... + + def build_distill_losses(self, losses): + ... + + def compute_distill_losses(self): + ... +``` + +4. Develop custom loss (Optional). + +Here we take `L1Loss` as an example. Create a new file in `mmrazor/models/losses/l1_loss.py`. + +```python +from mmrazor.registry import MODELS + +@MODELS.register_module() +class L1Loss(nn.Module): + def __init__( + self, + loss_weight: float = 1.0, + size_average: Optional[bool] = None, + reduce: Optional[bool] = None, + reduction: str = 'mean', + ) -> None: + super().__init__() + ... + + def forward(self, s_feature, t_feature): + loss = F.l1_loss(s_feature, t_feature, self.size_average, self.reduce, + self.reduction) + return self.loss_weight * loss +``` + +5. Import the class + +You can either add the following line to `mmrazor/models/algorithms/__init__.py` + +```Python +from .single_teacher_distill import SingleTeacherDistill + +__all__ = [..., 'SingleTeacherDistill'] +``` + +or alternatively add + +```Python +custom_imports = dict( + imports=['mmrazor.models.algorithms.distill.configurable.single_teacher_distill'], + allow_failed_imports=False) +``` + +to the config file to avoid modifying the original code. + +6. Use the algorithm in your config file + +```Python +algorithm = dict( + type='Distill', + distiller=dict(type='SingleTeacherDistill', ...), + # you can also use your new algorithm components here + ... +) +``` diff --git a/docs/en/user_guides/nas.md b/docs/en/user_guides/nas.md index 5a6d5e15f..ce05e5773 100644 --- a/docs/en/user_guides/nas.md +++ b/docs/en/user_guides/nas.md @@ -1 +1,136 @@ # Neural Architecture Search + +Here we show how to develop new NAS algorithms with an example of SPOS. + +1. Register a new algorithm + +Create a new file `mmrazor/models/algorithms/nas/spos.py`, class `SPOS` inherits from class `BaseAlgorithm` + +```Python +from mmrazor.registry import MODELS +from ..base import BaseAlgorithm + +@MODELS.register_module() +class SPOS(BaseAlgorithm): + def __init__(self, **kwargs): + super(SPOS, self).__init__(**kwargs) + pass + + def loss(self, batch_inputs, data_samples): + pass +``` + +2. Develop new algorithm components (optional) + +SPOS can directly use class `OneShotModuleMutator` as core functions provider. If mutators provided in MMRazor don’t meet your needs, you can develop new algorithm components for your algorithm like `OneShotModuleMutator`, we will take `OneShotModuleMutator` as an example to introduce how to develop a new algorithm component: + +a. Create a new file `mmrazor/models/mutators/module_mutator/one_shot_module_mutator.py`, class `OneShotModuleMutator` inherits from class `ModuleMutator` + +b. Finish the functions you need in `OneShotModuleMutator`, eg: `sample_choices`, `set_choices` and so on. + +```Python +from mmrazor.registry import MODELS +from .module_mutator import ModuleMutator + + +@MODELS.register_module() +class OneShotModuleMutator(ModuleMutator): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + def sample_choices(self) -> Dict[int, Any]: + pass + + def set_choices(self, choices: Dict[int, Any]) -> None: + pass + + @property + def mutable_class_type(self): + return OneShotMutableModule +``` + +c. Import the new mutator + +You can either add the following line to `mmrazor/models/mutators/__init__.py` + +```Python +from .module_mutator import OneShotModuleMutator +``` + +or alternatively add + +```Python +custom_imports = dict( + imports=['mmrazor.models.mutators.module_mutator.one_shot_module_mutator'], + allow_failed_imports=False) +``` + +to the config file to avoid modifying the original code. + +d. Use the algorithm component in your config file + +```Python +mutator=dict(type='mmrazor.OneShotModuleMutator') +``` + +For further information, please refer to [Mutator ](https://aicarrier.feishu.cn/docx/doxcnmcie75HcbqkfBGaEoemBKg)for more details. + +3. Rewrite its `loss` function. + +Develop key logic of your algorithm in function`loss`. When having special steps to optimize, you should rewrite the function `train_step`. + +```Python +@MODELS.register_module() +class SPOS(BaseAlgorithm): + def __init__(self, **kwargs): + super(SPOS, self).__init__(**kwargs) + pass + + def sample_subnet(self): + pass + + def set_subnet(self, subnet): + pass + + def loss(self, batch_inputs, data_samples): + if self.is_supernet: + random_subnet = self.sample_subnet() + self.set_subnet(random_subnet) + return self.architecture(batch_inputs, data_samples, mode='loss') + else: + return self.architecture(batch_inputs, data_samples, mode='loss') +``` + +4. Add your custom functions (optional) + +After finishing your key logic in function `loss`, if you also need other custom functions, you can add them in class `SPOS` as follows. + +5. Import the class + +You can either add the following line to `mmrazor/models/algorithms/nas/__init__.py` + +```Python +from .spos import SPOS + +__all__ = ['SPOS'] +``` + +or alternatively add + +```Python +custom_imports = dict( + imports=['mmrazor.models.algorithms.nas.spos'], + allow_failed_imports=False) +``` + +to the config file to avoid modifying the original code. + +6. Use the algorithm in your config file + +```Python +model = dict( + type='mmrazor.SPOS', + architecture=supernet, + mutator=dict(type='mmrazor.OneShotModuleMutator')) +``` diff --git a/docs/en/user_guides/pruning.md b/docs/en/user_guides/pruning.md index 8bed33b3f..0d39e92f6 100644 --- a/docs/en/user_guides/pruning.md +++ b/docs/en/user_guides/pruning.md @@ -1 +1,152 @@ # Pruning + +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/prunning/autoslim.py`, class `AutoSlim` inherits from class `BaseAlgorithm`. + +```Python +from mmrazor.registry import MODELS +from .base import BaseAlgorithm + +@MODELS.register_module() +class AutoSlim(BaseAlgorithm): + def __init__(self, + mutator, + distiller, + architecture, + data_preprocessor, + init_cfg = None, + num_samples = 2) -> None: + super().__init__(**kwargs) + pass + + def train_step(self, data, optimizer): + pass +``` + +2. Develop new algorithm components (optional) + +AutoSlim can directly use class `OneShotChannelMutator` as core functions provider. If it can not meet your needs, you can develop new algorithm components for your algorithm like `OneShotChannalMutator`. We will take `OneShotChannelMutator` as an example to introduce how to develop a new algorithm component: + +a. Create a new file `mmrazor/models/mutators/channel_mutator/one_shot_channel_mutator.py`, class `OneShotChannelMutator` can inherits from `ChannelMutator`. + +b. Finish the functions you need, eg: `build_search_groups`, `set_choices` , `sample_choices` and so on + +```Python +from mmrazor.registry import MODELS +from .channel_mutator import ChannelMutator + + +@MODELS.register_module() +class OneShotChannelMutator(ChannelMutator): + def __init__(self, **kwargs): + super().__init__(**kwargs) + + def sample_choices(self): + pass + + def set_choices(self, choice_dict): + pass + + # supernet is a kind of architecture in `mmrazor/models/architectures/` + def build_search_groups(self, supernet): + pass +``` + +c. Import the module in `mmrazor/models/mutators/channel_mutator/__init__.py` + +```Python +from .one_shot_channel_mutator import OneShotChannelMutator + + __all__ = [..., 'OneShotChannelMutator'] +``` + +3. Rewrite its train_step + +Develop key logic of your algorithm in function`train_step` + +```Python +from mmrazor.registry import MODELS +from ..base import BaseAlgorithm + +@ALGORITHMS.register_module() +class AutoSlim(BaseAlgorithm): + def __init__(self, + mutator, + distiller, + architecture, + data_preprocessor, + init_cfg = None, + num_samples = 2) -> None: + super(AutoSlim, self).__init__(**kwargs) + pass + + def train_step(self, data: List[dict], + optim_wrapper: OptimWrapper) -> Dict[str, torch.Tensor]: + + def distill_step( + batch_inputs: torch.Tensor, data_samples: List[BaseDataElement] + ) -> Dict[str, torch.Tensor]: + ... + return subnet_losses + + batch_inputs, data_samples = self.data_preprocessor(data, True) + + total_losses = dict() + self.set_max_subnet() + with optim_wrapper.optim_context( + self), self.distiller.teacher_recorders: # type: ignore + max_subnet_losses = self(batch_inputs, data_samples, mode='loss') + parsed_max_subnet_losses, _ = self.parse_losses(max_subnet_losses) + optim_wrapper.update_params(parsed_max_subnet_losses) + total_losses.update(add_prefix(max_subnet_losses, 'max_subnet')) + + self.set_min_subnet() + min_subnet_losses = distill_step(batch_inputs, data_samples) + total_losses.update(add_prefix(min_subnet_losses, 'min_subnet')) + + for sample_idx in range(self.num_samples): + self.set_subnet(self.sample_subnet()) + random_subnet_losses = distill_step(batch_inputs, data_samples) + total_losses.update( + add_prefix(random_subnet_losses, + f'random_subnet_{sample_idx}')) + + return total_losses +``` + +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` + +```Python +from .pruning import AutoSlim + +__all__ = [..., 'AutoSlim'] +``` + +Or alternatively add + +```Python +custom_imports = dict( + imports=['mmrazor.models.algorithms.pruning.autoslim'], + allow_failed_imports=False) +``` + +to the config file to avoid modifying the original code. + +6. Use the algorithm in your config file + +```Python +model = dict( + type='AutoSlim', + architecture=..., + mutator=dict(type='OneShotChannelMutator', ...), + ) +```