[Model Compression] auto compression

examples/model_compress/auto_compress/torch/auto_compress_module.py

@@ -0,0 +1,120 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+from typing import Callable, Optional
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+from torch.optim.lr_scheduler import StepLR
+from torchvision import datasets, transforms
+
+from nni.algorithms.compression.pytorch.auto_compress import AbstractAutoCompressModule
+
+torch.manual_seed(1)
+
+class LeNet(nn.Module):
+    def __init__(self):
+        super(LeNet, self).__init__()
+        self.conv1 = nn.Conv2d(1, 32, 3, 1)
+        self.conv2 = nn.Conv2d(32, 64, 3, 1)
+        self.dropout1 = nn.Dropout2d(0.25)
+        self.dropout2 = nn.Dropout2d(0.5)
+        self.fc1 = nn.Linear(9216, 128)
+        self.fc2 = nn.Linear(128, 10)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = F.relu(x)
+        x = self.conv2(x)
+        x = F.relu(x)
+        x = F.max_pool2d(x, 2)
+        x = self.dropout1(x)
+        x = torch.flatten(x, 1)
+        x = self.fc1(x)
+        x = F.relu(x)
+        x = self.dropout2(x)
+        x = self.fc2(x)
+        output = F.log_softmax(x, dim=1)
+        return output
+
+_use_cuda = torch.cuda.is_available()
+
+_train_kwargs = {'batch_size': 64}
+_test_kwargs = {'batch_size': 1000}
+if _use_cuda:
+    _cuda_kwargs = {'num_workers': 1,
+                    'pin_memory': True,
+                    'shuffle': True}
+    _train_kwargs.update(_cuda_kwargs)
+    _test_kwargs.update(_cuda_kwargs)
+
+_transform = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize((0.1307,), (0.3081,))
+])
+
+_dataset1 = datasets.MNIST('./data', train=True, download=True, transform=_transform)
+_dataset2 = datasets.MNIST('./data', train=False, transform=_transform)
+_train_loader = torch.utils.data.DataLoader(_dataset1, **_train_kwargs)
+_test_loader = torch.utils.data.DataLoader(_dataset2, **_test_kwargs)
+
+_device = torch.device("cuda" if _use_cuda else "cpu")
+_epoch = 2
+
+def _train(model, optimizer):
+    model.train()
+    for data, target in _train_loader:
+        data, target = data.to(_device), target.to(_device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+
+def _test(model):
+    model.eval()
+    test_loss = 0
+    correct = 0
+    with torch.no_grad():
+        for data, target in _test_loader:
+            data, target = data.to(_device), target.to(_device)
+            output = model(data)
+            test_loss += F.nll_loss(output, target, reduction='sum').item()
+            pred = output.argmax(dim=1, keepdim=True)
+            correct += pred.eq(target.view_as(pred)).sum().item()
+    test_loss /= len(_test_loader.dataset)
+    acc = 100 * correct / len(_test_loader.dataset)
+    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
+        test_loss, correct, len(_test_loader.dataset), acc))
+    return acc
+
+_model = LeNet().to(_device)
+
+_pre_train_optimizer = optim.Adadelta(_model.parameters(), lr=1)
+_scheduler = StepLR(_pre_train_optimizer, step_size=1, gamma=0.7)
+for _ in range(_epoch):
+    _train(_model, _pre_train_optimizer)
+    _test(_model)
+    _scheduler.step()
+
+class AutoCompressModule(AbstractAutoCompressModule):
+    @classmethod
+    def model(cls) -> nn.Module:
+        return _model
+
+    @classmethod
+    def optimizer(cls) -> torch.optim.Optimizer:
+        return torch.optim.SGD(_model.parameters(), lr=0.01)
+
+    @classmethod
+    def evaluator(cls) -> Callable[[nn.Module], float]:
+        return _test
+
+    @classmethod
+    def finetune_trainer(cls, compressor_type: str, algorithm_name: str) -> Optional[Callable[[nn.Module, optim.Optimizer], None]]:
+        def _trainer(model, optimizer):
+            for _ in range(_epoch):
+                _train(model, optimizer)
+        return _trainer

examples/model_compress/auto_compress/torch/auto_compress_torch.py

@@ -0,0 +1,48 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+from pathlib import Path
+
+from nni.algorithms.compression.pytorch.auto_compress import AutoCompressExperiment, AutoCompressSearchSpaceGenerator
+
+generator = AutoCompressSearchSpaceGenerator()
+generator.add_pruner_config('level', [
+    {
+        "sparsity": {
+            "_type": "uniform",
+            "_value": [0.01, 0.99]
+        },
+        'op_types': ['default']
+    }
+])
+generator.add_pruner_config('l1', [
+    {
+        "sparsity": {
+            "_type": "uniform",
+            "_value": [0.01, 0.99]
+        },
+        'op_types': ['Conv2d']
+    }
+])
+generator.add_quantizer_config('qat', [
+    {
+        'quant_types': ['weight', 'output'],
+        'quant_bits': {
+            'weight': 8,
+            'output': 8
+        },
+        'op_types': ['Conv2d', 'Linear']
+    }])
+search_space = generator.dumps()
+
+experiment = AutoCompressExperiment('local')
+experiment.config.experiment_name = 'auto compress torch example'
+experiment.config.trial_concurrency = 1
+experiment.config.max_trial_number = 10
+experiment.config.search_space = search_space
+experiment.config.trial_code_directory = Path(__file__).parent
+experiment.config.tuner.name = 'TPE'
+experiment.config.tuner.class_args['optimize_mode'] = 'maximize'
+experiment.config.training_service.use_active_gpu = True
+
+experiment.run(8080)

nni/algorithms/compression/pytorch/auto_compress/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+from .experiment import AutoCompressExperimentConfig, AutoCompressExperiment
+from .interface import AbstractAutoCompressModule
+from .utils import AutoCompressSearchSpaceGenerator

nni/algorithms/compression/pytorch/auto_compress/auto_compress_engine.py

@@ -0,0 +1,150 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import logging
+import json_tricks
+from typing import Optional, Callable
+
+from torch.nn import Module
+from torch.optim import Optimizer
+
+import nni
+from nni.retiarii.utils import import_
+from .constants import PRUNER_DICT, QUANTIZER_DICT
+from .interface import AbstractExecutionEngine
+
+_logger = logging.getLogger(__name__)
+_logger.setLevel(logging.INFO)
+
+class AutoCompressEngine(AbstractExecutionEngine):
+    @classmethod
+    def __convert_pruner_config_list(cls, converted_config_dict: dict) -> list:
+        config_dict = {}
+        for key, value in converted_config_dict.items():
+            _, op_types, op_names, var_name = key.split('::')
+            config_dict.setdefault((op_types, op_names), {})
+            config_dict[(op_types, op_names)][var_name] = value
+
+        config_list = []
+        for key, config in config_dict.items():
+            op_types, op_names = key
+            op_types = op_types.split(':') if op_types else []
+            op_names = op_names.split(':') if op_names else []
+            if op_types:
+                config['op_types'] = op_types
+            if op_names:
+                config['op_names'] = op_names
+            if 'op_types' in config or 'op_names' in config:
+                config_list.append(config)
+
+        return config_list
+
+    @classmethod
+    def __convert_quantizer_config_list(cls, converted_config_dict: dict) -> list:
+        config_dict = {}
+        for key, value in converted_config_dict.items():
+            _, quant_types, op_types, op_names, var_name = key.split('::')
+            config_dict.setdefault((quant_types, op_types, op_names), {})
+            config_dict[(quant_types, op_types, op_names)][var_name] = value
+
+        config_list = []
+        for key, config in config_dict.items():
+            quant_types, op_types, op_names = key
+            quant_types = quant_types.split(':')
+            op_types = op_types.split(':')
+            op_names = op_names.split(':')
+            if quant_types:
+                config['quant_types'] = quant_types
+            else:
+                continue
+            if op_types:
+                config['op_types'] = op_types
+            if op_names:
+                config['op_names'] = op_names
+            if 'op_types' in config or 'op_names' in config:
+                config_list.append(config)
+
+        return config_list
+
+    @classmethod
+    def _convert_config_list(cls, compressor_type: str, converted_config_dict: dict) -> list:
+        func_dict = {
+            'pruner': cls.__convert_pruner_config_list,
+            'quantizer': cls.__convert_quantizer_config_list
+        }
+        return func_dict[compressor_type](converted_config_dict)
+
+    @classmethod
+    def __compress_pruning_pipeline(cls, algorithm_name: str,
+                                    model: Module,
+                                    config_list: list,
+                                    evaluator: Callable[[Module], float],
+                                    optimizer: Optional[Optimizer],
+                                    trainer: Optional[Callable[[Module, Optimizer], None]],
+                                    finetune_trainer: Optional[Callable[[Module, Optimizer], None]],
+                                    **compressor_parameter_dict) -> Module:
+        # evaluator is for future use
+        pruner = PRUNER_DICT[algorithm_name](model, config_list, optimizer, **compressor_parameter_dict)
+        model = pruner.compress()
+        if trainer:
+            trainer(model)
+        if finetune_trainer:
+            finetune_trainer(model)
+        return model
+
+    @classmethod
+    def __compress_quantization_pipeline(cls, algorithm_name: str,
+                                         model: Module,
+                                         config_list: list,
+                                         evaluator: Callable[[Module], float],
+                                         optimizer: Optional[Optimizer],
+                                         trainer: Callable[[Module, Optimizer], None],
+                                         finetune_trainer: Optional[Callable[[Module, Optimizer], None]],
+                                         **compressor_parameter_dict) -> Module:
+        # evaluator is for future use
+        quantizer = QUANTIZER_DICT[algorithm_name](model, config_list, optimizer, **compressor_parameter_dict)
+        model = quantizer.compress()
+        if trainer:
+            trainer(model)
+        if finetune_trainer:
+            finetune_trainer(model)
+        return model
+
+    @classmethod
+    def _compress_pipeline(cls, compressor_type: str,
+                           algorithm_name: str,
+                           model: Module,
+                           config_list: list,
+                           evaluator: Callable[[Module], float],
+                           optimizer: Optional[Optimizer],
+                           trainer: Optional[Callable[[Module, Optimizer], None]],
+                           finetune_trainer: Optional[Callable[[Module, Optimizer], None]],
+                           **compressor_parameter_dict) -> Module:
+        func_dict = {
+            'pruner': cls.__compress_pruning_pipeline,
+            'quantizer': cls.__compress_quantization_pipeline
+        }
+        _logger.info('%s compressor config_list:\n%s', algorithm_name, json_tricks.dumps(config_list, indent=4))
+        return func_dict[compressor_type](algorithm_name, model, config_list, evaluator, optimizer, trainer,
+                                          finetune_trainer, **compressor_parameter_dict)
+
+    @classmethod
+    def trial_execute_compress(cls):
+        auto_compress_module = import_('auto_compress_module.AutoCompressModule')
+
+        parameter = nni.get_next_parameter()['compressor_type']
+        compressor_type, algorithm_config = parameter['_name'], parameter['algorithm_name']
+        algorithm_name = algorithm_config['_name']
+        converted_config_dict = {k: v for k, v in algorithm_config.items() if k.startswith('config_list::')}
+        parameter_dict = {k.split('parameter::')[1]: v for k, v in algorithm_config.items() if k.startswith('parameter::')}
+
+        config_list = cls._convert_config_list(compressor_type, converted_config_dict)
+
+        model, evaluator, optimizer = auto_compress_module.model(), auto_compress_module.evaluator(), auto_compress_module.optimizer()
+        trainer = auto_compress_module.trainer(compressor_type, algorithm_name)
+        finetune_trainer = auto_compress_module.trainer(compressor_type, algorithm_name)
+
+        compressed_model = cls._compress_pipeline(compressor_type, algorithm_name, model, config_list, evaluator,
+                                                  optimizer, trainer, finetune_trainer, **parameter_dict)
+
+        nni.report_final_result(evaluator(compressed_model))

nni/algorithms/compression/pytorch/auto_compress/constants.py

@@ -0,0 +1,25 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+from ..pruning import LevelPruner, SlimPruner, L1FilterPruner, L2FilterPruner, FPGMPruner, TaylorFOWeightFilterPruner, \
+    ActivationAPoZRankFilterPruner, ActivationMeanRankFilterPruner
+from ..quantization.quantizers import NaiveQuantizer, QAT_Quantizer, DoReFaQuantizer, BNNQuantizer
+
+
+PRUNER_DICT = {
+    'level': LevelPruner,
+    'slim': SlimPruner,
+    'l1': L1FilterPruner,
+    'l2': L2FilterPruner,
+    'fpgm': FPGMPruner,
+    'taylorfo': TaylorFOWeightFilterPruner,
+    'apoz': ActivationAPoZRankFilterPruner,
+    'mean_activation': ActivationMeanRankFilterPruner
+}
+
+QUANTIZER_DICT = {
+    'naive': NaiveQuantizer,
+    'qat': QAT_Quantizer,
+    'dorefa': DoReFaQuantizer,
+    'bnn': BNNQuantizer
+}

nni/algorithms/compression/pytorch/auto_compress/experiment.py

@@ -0,0 +1,44 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import os
+
+import dataclasses
+from pathlib import Path
+from subprocess import Popen
+from typing import Optional
+
+import nni
+from nni.experiment import Experiment, ExperimentConfig, AlgorithmConfig
+
+
+class AutoCompressExperimentConfig(ExperimentConfig):
+    def __setattr__(self, key, value):
+        fixed_attrs = {'trial_command': 'python3 -m nni.algorithms.compression.pytorch.auto_compress.trial_entry'}
+        if key in fixed_attrs and type(value) is not type(dataclasses.MISSING) and fixed_attrs[key] != value:
+            raise AttributeError(f'{key} is not supposed to be set in AutoCompress mode by users!')
+        # 'trial_code_directory' is handled differently because the path will be converted to absolute path by us
+        if key == 'trial_code_directory' and not (value == Path('.') or os.path.isabs(value)):
+            raise AttributeError(f'{key} is not supposed to be set in AutoCompress mode by users!')
+        self.__dict__[key] = value
+
+class AutoCompressExperiment(Experiment):
+    def __init__(self, config=None, training_service=None):
+        nni.runtime.log.init_logger_experiment()
+
+        self.config: Optional[ExperimentConfig] = None
+        self.id: Optional[str] = None
+        self.port: Optional[int] = None
+        self._proc: Optional[Popen] = None
+        self.mode = 'new'
+
+        args = [config, training_service]  # deal with overloading
+        if isinstance(args[0], (str, list)):
+            self.config = AutoCompressExperimentConfig(args[0])
+            self.config.tuner = AlgorithmConfig(name='_none_', class_args={})
+            self.config.assessor = AlgorithmConfig(name='_none_', class_args={})
+            self.config.advisor = AlgorithmConfig(name='_none_', class_args={})
+        else:
+            self.config = args[0]
+
+        self.config.trial_command = 'python3 -m nni.algorithms.compression.pytorch.auto_compress.trial_entry'