From 01f22955bfd83b13c75e9a98fa4f2796b541a6d5 Mon Sep 17 00:00:00 2001 From: MohamedAliRashad Date: Fri, 19 Feb 2021 15:23:30 +0200 Subject: [PATCH 01/10] Abstract class for automatic compression --- nni/compression/pytorch/auto_compress.py | 44 ++++++++++++++++++++++++ 1 file changed, 44 insertions(+) create mode 100644 nni/compression/pytorch/auto_compress.py diff --git a/nni/compression/pytorch/auto_compress.py b/nni/compression/pytorch/auto_compress.py new file mode 100644 index 0000000000..0a5e789c65 --- /dev/null +++ b/nni/compression/pytorch/auto_compress.py @@ -0,0 +1,44 @@ +import logging +from torch.nn import Module + +_logger = logging.getLogger(__name__) + + +class AutoCompressor: + """ + Auto Compression Class to utilize pruning & quantization in one place + """ + + def __init__(self) -> None: + """ + Initialize required pruners and quantizers for automatic compressing + + Parameters + ---------- + pruners : list + pruners chosen for compression + quantizers : [list, string] + quantizers chosen for model speedup (can choose "All") + training_dataloader : pytorch dataloader + training dataset to use for model finetuning + testing_dataloader : pytorch dataloader + testing dataset to evaluate model + evaluation_metric : pytorch functional + metric to evaluate the compression result + optimizer: pytorch optimizer + optimizer used to train the model + """ + pass + + def __call__(self) -> Module: + """ + Compressing the model + + Parameters + ---------- + model : pytorch model + the model user wants to compress + verbose : int + level of transparency in logging + """ + pass \ No newline at end of file From f9b3de51f78cd98a15da7ba0622103d02166b063 Mon Sep 17 00:00:00 2001 From: J-shang Date: Thu, 15 Apr 2021 09:39:08 +0800 Subject: [PATCH 02/10] temporary sync --- .../pruning/multi_prune_torch.py | 156 ++++++++++++++++++ .../compression/pytorch/multicompressor.py | 122 ++++++++++++++ 2 files changed, 278 insertions(+) create mode 100644 examples/model_compress/pruning/multi_prune_torch.py create mode 100644 nni/algorithms/compression/pytorch/multicompressor.py diff --git a/examples/model_compress/pruning/multi_prune_torch.py b/examples/model_compress/pruning/multi_prune_torch.py new file mode 100644 index 0000000000..6f8f47d966 --- /dev/null +++ b/examples/model_compress/pruning/multi_prune_torch.py @@ -0,0 +1,156 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +''' +NNI example for quick start of pruning. +In this example, we use level pruner to prune the LeNet on MNIST. +''' + +import logging + +import argparse +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torchvision import datasets, transforms +from torch.optim.lr_scheduler import StepLR +from models.mnist.lenet import LeNet +from nni.algorithms.compression.pytorch.multicompressor import MultiCompressor + +import nni + +_logger = logging.getLogger('mnist_example') +_logger.setLevel(logging.INFO) + +def train(args, model, device, train_loader, optimizer, epoch): + model.train() + for batch_idx, (data, target) in enumerate(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() + if batch_idx % args.log_interval == 0: + print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( + epoch, batch_idx * len(data), len(train_loader.dataset), + 100. * batch_idx / len(train_loader), loss.item())) + if args.dry_run: + break + + +def test(model, device, test_loader): + 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 + +def main(args): + torch.manual_seed(args.seed) + use_cuda = not args.no_cuda and torch.cuda.is_available() + + device = torch.device("cuda" if use_cuda else "cpu") + + train_kwargs = {'batch_size': args.batch_size} + test_kwargs = {'batch_size': args.test_batch_size} + 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) + + model = LeNet().to(device) + optimizer = optim.Adadelta(model.parameters(), lr=args.lr) + + print('start pre-training') + scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma) + for epoch in range(1, args.epochs + 1): + train(args, model, device, train_loader, optimizer, epoch) + test(model, device, test_loader) + scheduler.step() + + torch.save(model.state_dict(), "pretrain_mnist_lenet.pt") + + # model.load_state_dict(torch.load("pretrain_mnist_lenet.pt")) + + print('start pruning') + optimizer_finetune = torch.optim.SGD(model.parameters(), lr=0.01) + + # create pruner + prune_config = [ + { + 'config_list': [{'sparsity': 0.9, 'op_types': ['default']}], + 'pruner': { + 'type': 'level', + 'args': {} + } + } + ] + + pruner = MultiCompressor(model, prune_config, optimizer_finetune) + model = pruner.compress() + + # fine-tuning + best_top1 = 0 + for epoch in range(1, args.epochs + 1): + train(args, model, device, train_loader, optimizer_finetune, epoch) + top1 = test(model, device, test_loader) + + if top1 > best_top1: + best_top1 = top1 + # Export the best model, 'model_path' stores state_dict of the pruned model, + # mask_path stores mask_dict of the pruned model + pruner.export_model(model_path='pruend_mnist_lenet.pt', mask_path='mask_mnist_lenet.pt') + +if __name__ == '__main__': + # Training settings + parser = argparse.ArgumentParser(description='PyTorch MNIST Example for model comporession') + parser.add_argument('--batch-size', type=int, default=64, metavar='N', + help='input batch size for training (default: 64)') + parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N', + help='input batch size for testing (default: 1000)') + parser.add_argument('--epochs', type=int, default=10, metavar='N', + help='number of epochs to train (default: 10)') + parser.add_argument('--lr', type=float, default=1.0, metavar='LR', + help='learning rate (default: 1.0)') + parser.add_argument('--gamma', type=float, default=0.7, metavar='M', + help='Learning rate step gamma (default: 0.7)') + parser.add_argument('--no-cuda', action='store_true', default=False, + help='disables CUDA training') + parser.add_argument('--dry-run', action='store_true', default=False, + help='quickly check a single pass') + parser.add_argument('--seed', type=int, default=1, metavar='S', + help='random seed (default: 1)') + parser.add_argument('--log-interval', type=int, default=10, metavar='N', + help='how many batches to wait before logging training status') + parser.add_argument('--sparsity', type=float, default=0.5, + help='target overall target sparsity') + args = parser.parse_args() + + main(args) diff --git a/nni/algorithms/compression/pytorch/multicompressor.py b/nni/algorithms/compression/pytorch/multicompressor.py new file mode 100644 index 0000000000..0dcde7b857 --- /dev/null +++ b/nni/algorithms/compression/pytorch/multicompressor.py @@ -0,0 +1,122 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import logging +from typing import Tuple, List, Optional, Union +import torch + +from nni.compression.pytorch.compressor import Compressor +from nni.algorithms.compression.pytorch.pruning.one_shot import LevelPruner, L1FilterPruner, L2FilterPruner +from nni.algorithms.compression.pytorch.quantization.quantizers import NaiveQuantizer, QAT_Quantizer + +_logger = logging.getLogger(__name__) + + +PRUNER_DICT = { + 'level': LevelPruner, + 'l1': L1FilterPruner, + 'l2': L2FilterPruner +} + +QUANTIZER_DICT = { + 'naive': NaiveQuantizer, + 'qat': QAT_Quantizer +} + +MIXED_CONFIGLIST_TEMP = [ + { + 'config_list': [{'sparsity': 0.9, 'op_types': ['default']}], + 'pruner': { + 'type': 'level', + 'args': {} + } + } +] + +CONCFIG_LIST_TYPE = List[Tuple[str, dict, list]] + + +class Trainer: + def __init__(self): + pass + + def run(self): + pass + + +class MultiCompressor: + def __init__(self, model, mixed_config_list: List[dict], optimizer=None, trainer=None): + self.bound_model = model + self.pruner_config_list, self.quantizer_config_list = self._convert_config_list(mixed_config_list) + self.optimizer = optimizer + self.trainer = trainer + + self.pruners = [] + self.quantizers = [] + + def _convert_config_list(self, mixed_config_list: List[dict]) -> Tuple[CONCFIG_LIST_TYPE, CONCFIG_LIST_TYPE]: + pruner_config_list = [] + quantizer_config_list = [] + for config in mixed_config_list: + if 'pruner' in config: + pruner = config.get('pruner') + pruner_config_list.append((pruner['type'], pruner['args'], config.get('config_list'))) + elif 'quantizer' in config: + quantizer = config.get('quantizer') + quantizer_config_list.append((pruner['type'], pruner['args'], config.get('config_list'))) + return pruner_config_list, quantizer_config_list + + def compress(self): + for pruner_name, pruner_args, config_list in self.pruner_config_list: + pruner = PRUNER_DICT[pruner_name](self.bound_model, config_list, self.optimizer, **pruner_args) + self.pruners.append(pruner) + self.bound_model = pruner.compress() + if self.trainer: + self.trainer.run() + for quantizer_name, quantizer_args, config_list in self.quantizer_config_list: + quantizer = QUANTIZER_DICT[quantizer_name](self.bound_model, config_list, self.optimizer, **pruner_args) + self.quantizers.append(quantizer) + self.bound_model = quantizer.compress() + if self.trainer: + self.trainer.run() + return self.bound_model + + def export_model(self, model_path: str, mask_path: str, onnx_path: str = None, + input_shape: Optional[Union[List, Tuple]] = None, device: torch.device = None): + assert model_path is not None, 'model_path must be specified' + mask_dict = {} + + for pruner in self.pruners: + pruner._unwrap_model() + for wrapper in pruner.get_modules_wrapper(): + weight_mask = wrapper.weight_mask + bias_mask = wrapper.bias_mask + if weight_mask is not None: + mask_sum = weight_mask.sum().item() + mask_num = weight_mask.numel() + _logger.debug('Layer: %s Sparsity: %.4f', wrapper.name, 1 - mask_sum / mask_num) + wrapper.module.weight.data = wrapper.module.weight.data.mul(weight_mask) + if bias_mask is not None: + wrapper.module.bias.data = wrapper.module.bias.data.mul(bias_mask) + # save mask to dict + mask_dict[wrapper.name] = {"weight": weight_mask, "bias": bias_mask} + + torch.save(self.bound_model.state_dict(), model_path) + _logger.info('Model state_dict saved to %s', model_path) + if mask_path is not None: + torch.save(mask_dict, mask_path) + _logger.info('Mask dict saved to %s', mask_path) + if onnx_path is not None: + assert input_shape is not None, 'input_shape must be specified to export onnx model' + # input info needed + if device is None: + device = torch.device('cpu') + input_data = torch.Tensor(*input_shape) + torch.onnx.export(self.bound_model, input_data.to(device), onnx_path) + _logger.info('Model in onnx with input shape %s saved to %s', input_data.shape, onnx_path) + + for pruner in self.pruners: + pruner._wrap_model() + +if __name__ == '__main__': + pass From f4dbf44d0c2cffdd055f62976c4491fa1e7d246e Mon Sep 17 00:00:00 2001 From: J-shang Date: Fri, 16 Apr 2021 14:23:54 +0800 Subject: [PATCH 03/10] temp sync --- .../pruning/multi_prune_torch.py | 153 +++++++++++------- .../compression/pytorch/multicompressor.py | 51 +++--- 2 files changed, 128 insertions(+), 76 deletions(-) diff --git a/examples/model_compress/pruning/multi_prune_torch.py b/examples/model_compress/pruning/multi_prune_torch.py index 6f8f47d966..005d8ee811 100644 --- a/examples/model_compress/pruning/multi_prune_torch.py +++ b/examples/model_compress/pruning/multi_prune_torch.py @@ -23,42 +23,68 @@ _logger = logging.getLogger('mnist_example') _logger.setLevel(logging.INFO) -def train(args, model, device, train_loader, optimizer, epoch): - model.train() - for batch_idx, (data, target) in enumerate(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() - if batch_idx % args.log_interval == 0: - print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( - epoch, batch_idx * len(data), len(train_loader.dataset), - 100. * batch_idx / len(train_loader), loss.item())) - if args.dry_run: - break - - -def test(model, device, test_loader): - 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) +class Trainer: + def __init__(self, device, train_loader, test_loader, epochs, log_interval=10): + self.device = device + self.train_loader = train_loader + self.test_loader = test_loader + self.epochs = epochs + + self.log_interval = log_interval + + def pretrain(self, model, optimizer): + print('start pre-training') + scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma) + for epoch in range(1, self.epochs + 1): + self.__train(model, optimizer, epoch) + self.__test(model) + scheduler.step() + + def finetune(self, model, optimizer, pruner): + best_top1 = 0 + for epoch in range(1, args.epochs + 1): + self.__train(model, optimizer, epoch) + top1 = self.__test(model) + + if top1 > best_top1: + best_top1 = top1 + # Export the best model, 'model_path' stores state_dict of the pruned model, + # mask_path stores mask_dict of the pruned model + pruner.export_model(model_path='pruend_mnist_lenet.pt', mask_path='mask_mnist_lenet.pt') + + def __train(self, model, optimizer, epoch): + model.train() + for batch_idx, (data, target) in enumerate(self.train_loader): + data, target = data.to(self.device), target.to(self.device) + optimizer.zero_grad() 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 + loss = F.nll_loss(output, target) + loss.backward() + optimizer.step() + if batch_idx % self.log_interval == 0: + print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( + epoch, batch_idx * len(data), len(self.train_loader.dataset), + 100. * batch_idx / len(self.train_loader), loss.item())) + + def __test(self, model): + model.eval() + test_loss = 0 + correct = 0 + with torch.no_grad(): + for data, target in self.test_loader: + data, target = data.to(self.device), target.to(self.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(self.test_loader.dataset) + acc = 100 * correct / len(self.test_loader.dataset) + + print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( + test_loss, correct, len(self.test_loader.dataset), acc)) + + return acc def main(args): torch.manual_seed(args.seed) @@ -85,49 +111,62 @@ def main(args): train_loader = torch.utils.data.DataLoader(dataset1, **train_kwargs) test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs) + epochs = args.epochs + log_interval = args.log_interval + + trainer = Trainer(device, train_loader, test_loader, epochs, log_interval) + model = LeNet().to(device) optimizer = optim.Adadelta(model.parameters(), lr=args.lr) - print('start pre-training') - scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma) - for epoch in range(1, args.epochs + 1): - train(args, model, device, train_loader, optimizer, epoch) - test(model, device, test_loader) - scheduler.step() + trainer.pretrain(model, optimizer) torch.save(model.state_dict(), "pretrain_mnist_lenet.pt") - # model.load_state_dict(torch.load("pretrain_mnist_lenet.pt")) - print('start pruning') optimizer_finetune = torch.optim.SGD(model.parameters(), lr=0.01) # create pruner + configure_list = [{ + 'quant_types': ['weight'], + 'quant_bits': { + 'weight': 8, + }, # you can just use `int` here because all `quan_types` share same bits length, see config for `ReLu6` below. + 'op_types': ['Conv2d', 'Linear'] + }, { + 'quant_types': ['output'], + 'quant_bits': 8, + 'quant_start_step': 1000, + 'op_types':['ReLU6'] + }] + prune_config = [ { - 'config_list': [{'sparsity': 0.9, 'op_types': ['default']}], + 'config_list': [{'sparsity': args.sparsity, 'op_types': ['Linear']}], 'pruner': { 'type': 'level', 'args': {} } + }, + { + 'config_list': [{'sparsity': args.sparsity, 'op_types': ['Conv2d']}], + 'pruner': { + 'type': 'l1', + 'args': {} + } + }, + { + 'config_list': configure_list, + 'quantizer': { + 'type': 'qat', + 'args': {} + } } ] - pruner = MultiCompressor(model, prune_config, optimizer_finetune) + pruner = MultiCompressor(model, prune_config, optimizer_finetune, trainer) model = pruner.compress() - # fine-tuning - best_top1 = 0 - for epoch in range(1, args.epochs + 1): - train(args, model, device, train_loader, optimizer_finetune, epoch) - top1 = test(model, device, test_loader) - - if top1 > best_top1: - best_top1 = top1 - # Export the best model, 'model_path' stores state_dict of the pruned model, - # mask_path stores mask_dict of the pruned model - pruner.export_model(model_path='pruend_mnist_lenet.pt', mask_path='mask_mnist_lenet.pt') - if __name__ == '__main__': # Training settings parser = argparse.ArgumentParser(description='PyTorch MNIST Example for model comporession') diff --git a/nni/algorithms/compression/pytorch/multicompressor.py b/nni/algorithms/compression/pytorch/multicompressor.py index 0dcde7b857..ab6cc224d3 100644 --- a/nni/algorithms/compression/pytorch/multicompressor.py +++ b/nni/algorithms/compression/pytorch/multicompressor.py @@ -23,16 +23,6 @@ 'qat': QAT_Quantizer } -MIXED_CONFIGLIST_TEMP = [ - { - 'config_list': [{'sparsity': 0.9, 'op_types': ['default']}], - 'pruner': { - 'type': 'level', - 'args': {} - } - } -] - CONCFIG_LIST_TYPE = List[Tuple[str, dict, list]] @@ -40,7 +30,7 @@ class Trainer: def __init__(self): pass - def run(self): + def finetune(self): pass @@ -63,7 +53,7 @@ def _convert_config_list(self, mixed_config_list: List[dict]) -> Tuple[CONCFIG_L pruner_config_list.append((pruner['type'], pruner['args'], config.get('config_list'))) elif 'quantizer' in config: quantizer = config.get('quantizer') - quantizer_config_list.append((pruner['type'], pruner['args'], config.get('config_list'))) + quantizer_config_list.append((quantizer['type'], quantizer['args'], config.get('config_list'))) return pruner_config_list, quantizer_config_list def compress(self): @@ -71,20 +61,21 @@ def compress(self): pruner = PRUNER_DICT[pruner_name](self.bound_model, config_list, self.optimizer, **pruner_args) self.pruners.append(pruner) self.bound_model = pruner.compress() - if self.trainer: - self.trainer.run() + if self.trainer and len(self.pruner_config_list) > 0: + self.trainer.finetune(self.bound_model, self.optimizer, self) for quantizer_name, quantizer_args, config_list in self.quantizer_config_list: - quantizer = QUANTIZER_DICT[quantizer_name](self.bound_model, config_list, self.optimizer, **pruner_args) + quantizer = QUANTIZER_DICT[quantizer_name](self.bound_model, config_list, self.optimizer, **quantizer_args) self.quantizers.append(quantizer) self.bound_model = quantizer.compress() - if self.trainer: - self.trainer.run() + if self.trainer and len(self.quantizer_config_list) > 0: + self.trainer.finetune(self.bound_model, self.optimizer, self) return self.bound_model - def export_model(self, model_path: str, mask_path: str, onnx_path: str = None, + def export_model(self, model_path: str, mask_path: str = None, calibration_path: str = None, onnx_path: str = None, input_shape: Optional[Union[List, Tuple]] = None, device: torch.device = None): assert model_path is not None, 'model_path must be specified' mask_dict = {} + calibration_config = {} for pruner in self.pruners: pruner._unwrap_model() @@ -94,18 +85,36 @@ def export_model(self, model_path: str, mask_path: str, onnx_path: str = None, if weight_mask is not None: mask_sum = weight_mask.sum().item() mask_num = weight_mask.numel() - _logger.debug('Layer: %s Sparsity: %.4f', wrapper.name, 1 - mask_sum / mask_num) + _logger.info('Layer: %s Sparsity: %.4f', wrapper.name, 1 - mask_sum / mask_num) wrapper.module.weight.data = wrapper.module.weight.data.mul(weight_mask) if bias_mask is not None: wrapper.module.bias.data = wrapper.module.bias.data.mul(bias_mask) # save mask to dict mask_dict[wrapper.name] = {"weight": weight_mask, "bias": bias_mask} + for quantizer in self.quantizers: + quantizer._unwrap_model() + for name, module in quantizer.bound_model.named_modules(): + if hasattr(module, 'weight_bit') or hasattr(module, 'activation_bit'): + calibration_config[name] = {} + if hasattr(module, 'weight_bit'): + calibration_config[name]['weight_bit'] = int(module.weight_bit) + calibration_config[name]['tracked_min_input'] = float(module.tracked_min_input) + calibration_config[name]['tracked_max_input'] = float(module.tracked_max_input) + if hasattr(module, 'activation_bit'): + calibration_config[name]['activation_bit'] = int(module.activation_bit) + calibration_config[name]['tracked_min_activation'] = float(module.tracked_min_activation) + calibration_config[name]['tracked_max_activation'] = float(module.tracked_max_activation) + quantizer._del_simulated_attr(module) + torch.save(self.bound_model.state_dict(), model_path) _logger.info('Model state_dict saved to %s', model_path) if mask_path is not None: torch.save(mask_dict, mask_path) _logger.info('Mask dict saved to %s', mask_path) + if calibration_path is not None: + torch.save(calibration_config, calibration_path) + _logger.info('Calibration config saved to %s', calibration_path) if onnx_path is not None: assert input_shape is not None, 'input_shape must be specified to export onnx model' # input info needed @@ -118,5 +127,9 @@ def export_model(self, model_path: str, mask_path: str, onnx_path: str = None, for pruner in self.pruners: pruner._wrap_model() + for quantizer in self.quantizers: + quantizer._wrap_model() + + if __name__ == '__main__': pass From a5a8f2b559a3ffeb8b438575be74740b783a178e Mon Sep 17 00:00:00 2001 From: J-shang Date: Fri, 16 Apr 2021 17:02:47 +0800 Subject: [PATCH 04/10] temp sync --- .../pruning/multi_prune_torch.py | 3 +- .../compression/pytorch/multicompressor.py | 114 ++++++++++++++++-- 2 files changed, 104 insertions(+), 13 deletions(-) diff --git a/examples/model_compress/pruning/multi_prune_torch.py b/examples/model_compress/pruning/multi_prune_torch.py index 005d8ee811..ee8a0ec288 100644 --- a/examples/model_compress/pruning/multi_prune_torch.py +++ b/examples/model_compress/pruning/multi_prune_torch.py @@ -165,7 +165,8 @@ def main(args): ] pruner = MultiCompressor(model, prune_config, optimizer_finetune, trainer) - model = pruner.compress() + model = pruner.compress(model_path='pruend_mnist_lenet.pt', mask_path='mask_mnist_lenet.pt', + calibration_path='calibration_mnist_lenet.pt', input_shape=[10, 1, 28, 28]) if __name__ == '__main__': # Training settings diff --git a/nni/algorithms/compression/pytorch/multicompressor.py b/nni/algorithms/compression/pytorch/multicompressor.py index ab6cc224d3..adfb9dac13 100644 --- a/nni/algorithms/compression/pytorch/multicompressor.py +++ b/nni/algorithms/compression/pytorch/multicompressor.py @@ -6,6 +6,7 @@ import torch from nni.compression.pytorch.compressor import Compressor +from nni.compression.pytorch.speedup import ModelSpeedup from nni.algorithms.compression.pytorch.pruning.one_shot import LevelPruner, L1FilterPruner, L2FilterPruner from nni.algorithms.compression.pytorch.quantization.quantizers import NaiveQuantizer, QAT_Quantizer @@ -56,26 +57,36 @@ def _convert_config_list(self, mixed_config_list: List[dict]) -> Tuple[CONCFIG_L quantizer_config_list.append((quantizer['type'], quantizer['args'], config.get('config_list'))) return pruner_config_list, quantizer_config_list - def compress(self): + def compress(self, model_path: str, mask_path: str = None, calibration_path: str = None, onnx_path: str = None, + input_shape: Optional[Union[List, Tuple]] = None): for pruner_name, pruner_args, config_list in self.pruner_config_list: pruner = PRUNER_DICT[pruner_name](self.bound_model, config_list, self.optimizer, **pruner_args) self.pruners.append(pruner) self.bound_model = pruner.compress() - if self.trainer and len(self.pruner_config_list) > 0: - self.trainer.finetune(self.bound_model, self.optimizer, self) + + if len(self.pruner_config_list) > 0: + mask_dict = self.export_pruned_model(model_path, mask_path, onnx_path, input_shape, self.device) + dummy_input = torch.randn(input_shape).to(self.device) + model_sp = ModelSpeedup(self.bound_model, dummy_input, mask_dict) + model_sp.speedup_model() + + if self.trainer: + self.trainer.finetune(self.bound_model, self.optimizer, self) + for quantizer_name, quantizer_args, config_list in self.quantizer_config_list: quantizer = QUANTIZER_DICT[quantizer_name](self.bound_model, config_list, self.optimizer, **quantizer_args) self.quantizers.append(quantizer) self.bound_model = quantizer.compress() + if self.trainer and len(self.quantizer_config_list) > 0: self.trainer.finetune(self.bound_model, self.optimizer, self) + calibration_config = self.export_quantized_model(model_path, calibration_path, onnx_path, input_shape, self.device) return self.bound_model - def export_model(self, model_path: str, mask_path: str = None, calibration_path: str = None, onnx_path: str = None, - input_shape: Optional[Union[List, Tuple]] = None, device: torch.device = None): + def export_pruned_model(self, model_path: str, mask_path: str = None, onnx_path: str = None, + input_shape: Optional[Union[List, Tuple]] = None, device: torch.device = None): assert model_path is not None, 'model_path must be specified' mask_dict = {} - calibration_config = {} for pruner in self.pruners: pruner._unwrap_model() @@ -92,6 +103,30 @@ def export_model(self, model_path: str, mask_path: str = None, calibration_path: # save mask to dict mask_dict[wrapper.name] = {"weight": weight_mask, "bias": bias_mask} + torch.save(self.bound_model.state_dict(), model_path) + _logger.info('Model state_dict saved to %s', model_path) + if mask_path is not None: + torch.save(mask_dict, mask_path) + _logger.info('Mask dict saved to %s', mask_path) + if onnx_path is not None: + assert input_shape is not None, 'input_shape must be specified to export onnx model' + # input info needed + if device is None: + device = torch.device('cpu') + input_data = torch.Tensor(*input_shape) + torch.onnx.export(self.bound_model, input_data.to(device), onnx_path) + _logger.info('Model in onnx with input shape %s saved to %s', input_data.shape, onnx_path) + + for pruner in self.pruners: + pruner._wrap_model() + + return mask_dict + + def export_quantized_model(self, model_path: str, calibration_path: str = None, onnx_path: str = None, + input_shape: Optional[Union[List, Tuple]] = None, device: torch.device = None): + assert model_path is not None, 'model_path must be specified' + calibration_config = {} + for quantizer in self.quantizers: quantizer._unwrap_model() for name, module in quantizer.bound_model.named_modules(): @@ -109,9 +144,6 @@ def export_model(self, model_path: str, mask_path: str = None, calibration_path: torch.save(self.bound_model.state_dict(), model_path) _logger.info('Model state_dict saved to %s', model_path) - if mask_path is not None: - torch.save(mask_dict, mask_path) - _logger.info('Mask dict saved to %s', mask_path) if calibration_path is not None: torch.save(calibration_config, calibration_path) _logger.info('Calibration config saved to %s', calibration_path) @@ -124,12 +156,70 @@ def export_model(self, model_path: str, mask_path: str = None, calibration_path: torch.onnx.export(self.bound_model, input_data.to(device), onnx_path) _logger.info('Model in onnx with input shape %s saved to %s', input_data.shape, onnx_path) - for pruner in self.pruners: - pruner._wrap_model() - for quantizer in self.quantizers: quantizer._wrap_model() + return calibration_config + + # def export_model(self, model_path: str, mask_path: str = None, calibration_path: str = None, onnx_path: str = None, + # input_shape: Optional[Union[List, Tuple]] = None, device: torch.device = None): + # assert model_path is not None, 'model_path must be specified' + # mask_dict = {} + # calibration_config = {} + + # for pruner in self.pruners: + # pruner._unwrap_model() + # for wrapper in pruner.get_modules_wrapper(): + # weight_mask = wrapper.weight_mask + # bias_mask = wrapper.bias_mask + # if weight_mask is not None: + # mask_sum = weight_mask.sum().item() + # mask_num = weight_mask.numel() + # _logger.info('Layer: %s Sparsity: %.4f', wrapper.name, 1 - mask_sum / mask_num) + # wrapper.module.weight.data = wrapper.module.weight.data.mul(weight_mask) + # if bias_mask is not None: + # wrapper.module.bias.data = wrapper.module.bias.data.mul(bias_mask) + # # save mask to dict + # mask_dict[wrapper.name] = {"weight": weight_mask, "bias": bias_mask} + + # for quantizer in self.quantizers: + # quantizer._unwrap_model() + # for name, module in quantizer.bound_model.named_modules(): + # if hasattr(module, 'weight_bit') or hasattr(module, 'activation_bit'): + # calibration_config[name] = {} + # if hasattr(module, 'weight_bit'): + # calibration_config[name]['weight_bit'] = int(module.weight_bit) + # calibration_config[name]['tracked_min_input'] = float(module.tracked_min_input) + # calibration_config[name]['tracked_max_input'] = float(module.tracked_max_input) + # if hasattr(module, 'activation_bit'): + # calibration_config[name]['activation_bit'] = int(module.activation_bit) + # calibration_config[name]['tracked_min_activation'] = float(module.tracked_min_activation) + # calibration_config[name]['tracked_max_activation'] = float(module.tracked_max_activation) + # quantizer._del_simulated_attr(module) + + # torch.save(self.bound_model.state_dict(), model_path) + # _logger.info('Model state_dict saved to %s', model_path) + # if mask_path is not None: + # torch.save(mask_dict, mask_path) + # _logger.info('Mask dict saved to %s', mask_path) + # if calibration_path is not None: + # torch.save(calibration_config, calibration_path) + # _logger.info('Calibration config saved to %s', calibration_path) + # if onnx_path is not None: + # assert input_shape is not None, 'input_shape must be specified to export onnx model' + # # input info needed + # if device is None: + # device = torch.device('cpu') + # input_data = torch.Tensor(*input_shape) + # torch.onnx.export(self.bound_model, input_data.to(device), onnx_path) + # _logger.info('Model in onnx with input shape %s saved to %s', input_data.shape, onnx_path) + + # for pruner in self.pruners: + # pruner._wrap_model() + + # for quantizer in self.quantizers: + # quantizer._wrap_model() + if __name__ == '__main__': pass From 9641f26a684777ec95add995cdf73ffae1889200 Mon Sep 17 00:00:00 2001 From: J-shang Date: Tue, 20 Apr 2021 16:30:30 +0800 Subject: [PATCH 05/10] temp sync --- .../pruning/multi_prune_torch.py | 11 +- .../compression/pytorch/multicompressor.py | 108 +++++++++++------- .../pytorch/quantization/quantizers.py | 2 +- 3 files changed, 76 insertions(+), 45 deletions(-) diff --git a/examples/model_compress/pruning/multi_prune_torch.py b/examples/model_compress/pruning/multi_prune_torch.py index ee8a0ec288..2bfd5fdb41 100644 --- a/examples/model_compress/pruning/multi_prune_torch.py +++ b/examples/model_compress/pruning/multi_prune_torch.py @@ -48,9 +48,9 @@ def finetune(self, model, optimizer, pruner): if top1 > best_top1: best_top1 = top1 - # Export the best model, 'model_path' stores state_dict of the pruned model, - # mask_path stores mask_dict of the pruned model - pruner.export_model(model_path='pruend_mnist_lenet.pt', mask_path='mask_mnist_lenet.pt') + pruner.save_bound_model('bound_model.pt') + + return 'bound_model.pt' def __train(self, model, optimizer, epoch): model.train() @@ -165,8 +165,9 @@ def main(args): ] pruner = MultiCompressor(model, prune_config, optimizer_finetune, trainer) - model = pruner.compress(model_path='pruend_mnist_lenet.pt', mask_path='mask_mnist_lenet.pt', - calibration_path='calibration_mnist_lenet.pt', input_shape=[10, 1, 28, 28]) + pruner.set_config(model_path='pruend_mnist_lenet.pt', mask_path='mask_mnist_lenet.pt', + calibration_path='calibration_mnist_lenet.pt', input_shape=[10, 1, 28, 28], device=device) + model = pruner.compress() if __name__ == '__main__': # Training settings diff --git a/nni/algorithms/compression/pytorch/multicompressor.py b/nni/algorithms/compression/pytorch/multicompressor.py index adfb9dac13..42cfcce388 100644 --- a/nni/algorithms/compression/pytorch/multicompressor.py +++ b/nni/algorithms/compression/pytorch/multicompressor.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from enum import Enum import logging from typing import Tuple, List, Optional, Union import torch @@ -26,7 +27,6 @@ CONCFIG_LIST_TYPE = List[Tuple[str, dict, list]] - class Trainer: def __init__(self): pass @@ -45,6 +45,22 @@ def __init__(self, model, mixed_config_list: List[dict], optimizer=None, trainer self.pruners = [] self.quantizers = [] + self.model_path = None + self.mask_path = None + self.calibration_path = None + self.onnx_path = None + self.input_shape = None + self.device = None + + def set_config(self, model_path: str, mask_path: str = None, calibration_path: str = None, onnx_path: str = None, + input_shape: Optional[Union[List, Tuple]] = None, device: torch.device = None): + self.model_path = model_path + self.mask_path = mask_path + self.calibration_path = calibration_path + self.onnx_path = onnx_path + self.input_shape = input_shape + self.device = device + def _convert_config_list(self, mixed_config_list: List[dict]) -> Tuple[CONCFIG_LIST_TYPE, CONCFIG_LIST_TYPE]: pruner_config_list = [] quantizer_config_list = [] @@ -57,35 +73,50 @@ def _convert_config_list(self, mixed_config_list: List[dict]) -> Tuple[CONCFIG_L quantizer_config_list.append((quantizer['type'], quantizer['args'], config.get('config_list'))) return pruner_config_list, quantizer_config_list - def compress(self, model_path: str, mask_path: str = None, calibration_path: str = None, onnx_path: str = None, - input_shape: Optional[Union[List, Tuple]] = None): + def compress(self): for pruner_name, pruner_args, config_list in self.pruner_config_list: pruner = PRUNER_DICT[pruner_name](self.bound_model, config_list, self.optimizer, **pruner_args) self.pruners.append(pruner) + print('Use {} pruner, start pruning...'.format(pruner_name)) self.bound_model = pruner.compress() if len(self.pruner_config_list) > 0: - mask_dict = self.export_pruned_model(model_path, mask_path, onnx_path, input_shape, self.device) - dummy_input = torch.randn(input_shape).to(self.device) - model_sp = ModelSpeedup(self.bound_model, dummy_input, mask_dict) + self._export_pruned_model() + dummy_input = torch.randn(self.input_shape).to(self.device) + model_sp = ModelSpeedup(self.bound_model, dummy_input, self.mask_path) model_sp.speedup_model() + for pruner in self.pruners: + pruner._unwrap_model() + if self.trainer: - self.trainer.finetune(self.bound_model, self.optimizer, self) + saved_path = self.trainer.finetune(self.bound_model, self.optimizer, self) + if saved_path is not None: + self.bound_model = torch.load(saved_path) for quantizer_name, quantizer_args, config_list in self.quantizer_config_list: quantizer = QUANTIZER_DICT[quantizer_name](self.bound_model, config_list, self.optimizer, **quantizer_args) self.quantizers.append(quantizer) self.bound_model = quantizer.compress() + self.bound_model.to(self.device) + if self.trainer and len(self.quantizer_config_list) > 0: - self.trainer.finetune(self.bound_model, self.optimizer, self) - calibration_config = self.export_quantized_model(model_path, calibration_path, onnx_path, input_shape, self.device) + saved_path = self.trainer.finetune(self.bound_model, self.optimizer, self) + if saved_path is not None: + self.bound_model = torch.load(saved_path) + calibration_config = self._export_quantized_model() + return self.bound_model - def export_pruned_model(self, model_path: str, mask_path: str = None, onnx_path: str = None, - input_shape: Optional[Union[List, Tuple]] = None, device: torch.device = None): - assert model_path is not None, 'model_path must be specified' + def save_bound_model(self, path: str): + torch.save(self.bound_model, path) + + def load_bound_model(self, path: str): + self.bound_model = torch.load(path) + + def _export_pruned_model(self): + assert self.model_path is not None, 'model_path must be specified' mask_dict = {} for pruner in self.pruners: @@ -103,28 +134,27 @@ def export_pruned_model(self, model_path: str, mask_path: str = None, onnx_path: # save mask to dict mask_dict[wrapper.name] = {"weight": weight_mask, "bias": bias_mask} - torch.save(self.bound_model.state_dict(), model_path) - _logger.info('Model state_dict saved to %s', model_path) - if mask_path is not None: - torch.save(mask_dict, mask_path) - _logger.info('Mask dict saved to %s', mask_path) - if onnx_path is not None: - assert input_shape is not None, 'input_shape must be specified to export onnx model' + torch.save(self.bound_model.state_dict(), self.model_path) + _logger.info('Model state_dict saved to %s', self.model_path) + if self.mask_path is not None: + torch.save(mask_dict, self.mask_path) + _logger.info('Mask dict saved to %s', self.mask_path) + if self.onnx_path is not None: + assert self.input_shape is not None, 'input_shape must be specified to export onnx model' # input info needed - if device is None: - device = torch.device('cpu') - input_data = torch.Tensor(*input_shape) - torch.onnx.export(self.bound_model, input_data.to(device), onnx_path) - _logger.info('Model in onnx with input shape %s saved to %s', input_data.shape, onnx_path) + if self.device is None: + self.device = torch.device('cpu') + input_data = torch.Tensor(*self.input_shape) + torch.onnx.export(self.bound_model, input_data.to(self.device), self.onnx_path) + _logger.info('Model in onnx with input shape %s saved to %s', input_data.shape, self.onnx_path) for pruner in self.pruners: pruner._wrap_model() return mask_dict - def export_quantized_model(self, model_path: str, calibration_path: str = None, onnx_path: str = None, - input_shape: Optional[Union[List, Tuple]] = None, device: torch.device = None): - assert model_path is not None, 'model_path must be specified' + def _export_quantized_model(self): + assert self.model_path is not None, 'model_path must be specified' calibration_config = {} for quantizer in self.quantizers: @@ -142,19 +172,19 @@ def export_quantized_model(self, model_path: str, calibration_path: str = None, calibration_config[name]['tracked_max_activation'] = float(module.tracked_max_activation) quantizer._del_simulated_attr(module) - torch.save(self.bound_model.state_dict(), model_path) - _logger.info('Model state_dict saved to %s', model_path) - if calibration_path is not None: - torch.save(calibration_config, calibration_path) - _logger.info('Calibration config saved to %s', calibration_path) - if onnx_path is not None: - assert input_shape is not None, 'input_shape must be specified to export onnx model' + torch.save(self.bound_model.state_dict(), self.model_path) + _logger.info('Model state_dict saved to %s', self.model_path) + if self.calibration_path is not None: + torch.save(calibration_config, self.calibration_path) + _logger.info('Calibration config saved to %s', self.calibration_path) + if self.onnx_path is not None: + assert self.input_shape is not None, 'input_shape must be specified to export onnx model' # input info needed - if device is None: - device = torch.device('cpu') - input_data = torch.Tensor(*input_shape) - torch.onnx.export(self.bound_model, input_data.to(device), onnx_path) - _logger.info('Model in onnx with input shape %s saved to %s', input_data.shape, onnx_path) + if self.device is None: + self.device = torch.device('cpu') + input_data = torch.Tensor(*self.input_shape) + torch.onnx.export(self.bound_model, input_data.to(self.device), self.onnx_path) + _logger.info('Model in onnx with input shape %s saved to %s', input_data.shape, self.onnx_path) for quantizer in self.quantizers: quantizer._wrap_model() diff --git a/nni/algorithms/compression/pytorch/quantization/quantizers.py b/nni/algorithms/compression/pytorch/quantization/quantizers.py index ca40e30e45..e79e106f69 100644 --- a/nni/algorithms/compression/pytorch/quantization/quantizers.py +++ b/nni/algorithms/compression/pytorch/quantization/quantizers.py @@ -59,7 +59,7 @@ def update_ema(biased_ema, value, decay): float, float """ biased_ema = biased_ema * decay + (1 - decay) * value - return biased_ema + return biased_ema def update_quantization_param(bits, rmin, rmax): From cf401278f9294ff7b91c143ff8f0c8bb212b80e1 Mon Sep 17 00:00:00 2001 From: J-shang Date: Tue, 20 Apr 2021 16:54:05 +0800 Subject: [PATCH 06/10] temp sync --- nni/algorithms/compression/pytorch/multicompressor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nni/algorithms/compression/pytorch/multicompressor.py b/nni/algorithms/compression/pytorch/multicompressor.py index 42cfcce388..d6d34b59c8 100644 --- a/nni/algorithms/compression/pytorch/multicompressor.py +++ b/nni/algorithms/compression/pytorch/multicompressor.py @@ -105,7 +105,7 @@ def compress(self): saved_path = self.trainer.finetune(self.bound_model, self.optimizer, self) if saved_path is not None: self.bound_model = torch.load(saved_path) - calibration_config = self._export_quantized_model() + # calibration_config = self._export_quantized_model() return self.bound_model From 0d3b40e99d89e3e287329cddc1365958de3e42be Mon Sep 17 00:00:00 2001 From: J-shang Date: Wed, 21 Apr 2021 10:18:53 +0800 Subject: [PATCH 07/10] temp sync --- .../pruning/multi_prune_torch.py | 4 +-- .../compression/pytorch/multicompressor.py | 25 ++++++++----------- 2 files changed, 12 insertions(+), 17 deletions(-) diff --git a/examples/model_compress/pruning/multi_prune_torch.py b/examples/model_compress/pruning/multi_prune_torch.py index 2bfd5fdb41..0cc2e36982 100644 --- a/examples/model_compress/pruning/multi_prune_torch.py +++ b/examples/model_compress/pruning/multi_prune_torch.py @@ -48,9 +48,9 @@ def finetune(self, model, optimizer, pruner): if top1 > best_top1: best_top1 = top1 - pruner.save_bound_model('bound_model.pt') + pruner.save('multicompressor_dict.pkl') - return 'bound_model.pt' + return 'multicompressor_dict.pkl' def __train(self, model, optimizer, epoch): model.train() diff --git a/nni/algorithms/compression/pytorch/multicompressor.py b/nni/algorithms/compression/pytorch/multicompressor.py index d6d34b59c8..25b0954951 100644 --- a/nni/algorithms/compression/pytorch/multicompressor.py +++ b/nni/algorithms/compression/pytorch/multicompressor.py @@ -1,8 +1,8 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -from enum import Enum import logging +import pickle from typing import Tuple, List, Optional, Union import torch @@ -27,13 +27,6 @@ CONCFIG_LIST_TYPE = List[Tuple[str, dict, list]] -class Trainer: - def __init__(self): - pass - - def finetune(self): - pass - class MultiCompressor: def __init__(self, model, mixed_config_list: List[dict], optimizer=None, trainer=None): @@ -92,7 +85,7 @@ def compress(self): if self.trainer: saved_path = self.trainer.finetune(self.bound_model, self.optimizer, self) if saved_path is not None: - self.bound_model = torch.load(saved_path) + self.load(saved_path) for quantizer_name, quantizer_args, config_list in self.quantizer_config_list: quantizer = QUANTIZER_DICT[quantizer_name](self.bound_model, config_list, self.optimizer, **quantizer_args) @@ -104,16 +97,18 @@ def compress(self): if self.trainer and len(self.quantizer_config_list) > 0: saved_path = self.trainer.finetune(self.bound_model, self.optimizer, self) if saved_path is not None: - self.bound_model = torch.load(saved_path) - # calibration_config = self._export_quantized_model() + self.load(saved_path) + calibration_config = self._export_quantized_model() return self.bound_model - def save_bound_model(self, path: str): - torch.save(self.bound_model, path) + def save(self, path: str): + with open(path, 'wb') as f: + pickle.dump(self.__dict__, f) - def load_bound_model(self, path: str): - self.bound_model = torch.load(path) + def load(self, path: str): + with open(path, 'rb') as f: + self.__dict__ = pickle.load(f) def _export_pruned_model(self): assert self.model_path is not None, 'model_path must be specified' From 005356ec2f464960097f9b3620fcb3fe6373f33d Mon Sep 17 00:00:00 2001 From: J-shang Date: Wed, 21 Apr 2021 17:25:01 +0800 Subject: [PATCH 08/10] add more pruner --- .../compression/pytorch/multicompressor.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/nni/algorithms/compression/pytorch/multicompressor.py b/nni/algorithms/compression/pytorch/multicompressor.py index 25b0954951..7bb2f6b9d9 100644 --- a/nni/algorithms/compression/pytorch/multicompressor.py +++ b/nni/algorithms/compression/pytorch/multicompressor.py @@ -8,21 +8,29 @@ from nni.compression.pytorch.compressor import Compressor from nni.compression.pytorch.speedup import ModelSpeedup -from nni.algorithms.compression.pytorch.pruning.one_shot import LevelPruner, L1FilterPruner, L2FilterPruner -from nni.algorithms.compression.pytorch.quantization.quantizers import NaiveQuantizer, QAT_Quantizer +from nni.algorithms.compression.pytorch.pruning.one_shot import LevelPruner, SlimPruner, L1FilterPruner, L2FilterPruner, FPGMPruner, + TaylorFOWeightFilterPruner, ActivationAPoZRankFilterPruner, ActivationMeanRankFilterPruner +from nni.algorithms.compression.pytorch.quantization.quantizers import NaiveQuantizer, QAT_Quantizer, DoReFaQuantizer, BNNQuantizer _logger = logging.getLogger(__name__) PRUNER_DICT = { 'level': LevelPruner, + 'slim': SlimPruner, 'l1': L1FilterPruner, - 'l2': L2FilterPruner + 'l2': L2FilterPruner, + 'fpgm': FPGMPruner, + 'taylorfo': TaylorFOWeightFilterPruner, + 'apoz': ActivationAPoZRankFilterPruner, + 'mean_activation': ActivationMeanRankFilterPruner } QUANTIZER_DICT = { 'naive': NaiveQuantizer, - 'qat': QAT_Quantizer + 'qat': QAT_Quantizer, + 'dorefa': DoReFaQuantizer, + 'bnn': BNNQuantizer } CONCFIG_LIST_TYPE = List[Tuple[str, dict, list]] From 65089b06f5817466d3d53d5b23252ea641d7ec22 Mon Sep 17 00:00:00 2001 From: J-shang Date: Sun, 25 Apr 2021 16:59:28 +0800 Subject: [PATCH 09/10] temp sync --- .../model_compress/auto_compress_torch.py | 152 ++++++++++++++++++ ...une_torch.py => multi_compressor_torch.py} | 11 +- .../compression/pytorch/autocompressor.py | 108 +++++++++++++ .../compression/pytorch/multicompressor.py | 6 +- 4 files changed, 266 insertions(+), 11 deletions(-) create mode 100644 examples/model_compress/auto_compress_torch.py rename examples/model_compress/{pruning/multi_prune_torch.py => multi_compressor_torch.py} (97%) create mode 100644 nni/algorithms/compression/pytorch/autocompressor.py diff --git a/examples/model_compress/auto_compress_torch.py b/examples/model_compress/auto_compress_torch.py new file mode 100644 index 0000000000..82b242c8c0 --- /dev/null +++ b/examples/model_compress/auto_compress_torch.py @@ -0,0 +1,152 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import argparse +import logging +import torch + +from torchvision import datasets, transforms +import torch.nn.functional as F +import torch.optim as optim +from torch.optim.lr_scheduler import StepLR +from pruning.models.mnist.lenet import LeNet + +from nni.algorithms.compression.pytorch.autocompressor import AutoCompressor + +_logger = logging.getLogger(__name__) + +class Trainer: + def __init__(self, device, train_loader, test_loader, epochs, gamma, log_interval=10): + self.device = device + self.train_loader = train_loader + self.test_loader = test_loader + self.epochs = epochs + self.gamma = gamma + + self.log_interval = log_interval + + def pretrain(self, model, optimizer): + print('start pre-training') + scheduler = StepLR(optimizer, step_size=1, gamma=self.gamma) + for epoch in range(1, self.epochs + 1): + self.__train(model, optimizer, epoch) + self.__test(model) + scheduler.step() + + def finetune(self, model, optimizer, pruner): + best_top1 = 0 + for epoch in range(1, self.epochs + 1): + self.__train(model, optimizer, epoch) + top1 = self.__test(model) + + if top1 > best_top1: + best_top1 = top1 + pruner.save('multicompressor_dict.pkl') + + return 'multicompressor_dict.pkl' + + def __train(self, model, optimizer, epoch): + model.train() + for batch_idx, (data, target) in enumerate(self.train_loader): + data, target = data.to(self.device), target.to(self.device) + optimizer.zero_grad() + output = model(data) + loss = F.nll_loss(output, target) + loss.backward() + optimizer.step() + if batch_idx % self.log_interval == 0: + print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( + epoch, batch_idx * len(data), len(self.train_loader.dataset), + 100. * batch_idx / len(self.train_loader), loss.item())) + + def __test(self, model): + model.eval() + test_loss = 0 + correct = 0 + with torch.no_grad(): + for data, target in self.test_loader: + data, target = data.to(self.device), target.to(self.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(self.test_loader.dataset) + acc = 100 * correct / len(self.test_loader.dataset) + + print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( + test_loss, correct, len(self.test_loader.dataset), acc)) + + return acc + + +def main(args): + torch.manual_seed(args.seed) + use_cuda = not args.no_cuda and torch.cuda.is_available() + + device = torch.device("cuda" if use_cuda else "cpu") + + train_kwargs = {'batch_size': args.batch_size} + test_kwargs = {'batch_size': args.test_batch_size} + 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) + + epochs = args.epochs + log_interval = args.log_interval + + trainer = Trainer(device, train_loader, test_loader, epochs, log_interval) + + model = LeNet().to(device) + optimizer = optim.Adadelta(model.parameters(), lr=args.lr) + + trainer.pretrain(model, optimizer) + + torch.save(model.state_dict(), "pretrain_mnist_lenet.pt") + + print('start pruning') + optimizer_finetune = torch.optim.SGD(model.parameters(), lr=0.01) + + pruner = AutoCompressor(model, optimizer=optimizer_finetune, trainer=trainer) + + pruner.run(input_shape=[10, 1, 28, 28], device=device) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='PyTorch MNIST Example for model comporession') + parser.add_argument('--batch-size', type=int, default=64, metavar='N', + help='input batch size for training (default: 64)') + parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N', + help='input batch size for testing (default: 1000)') + parser.add_argument('--epochs', type=int, default=1, metavar='N', + help='number of epochs to train (default: 10)') + parser.add_argument('--lr', type=float, default=1.0, metavar='LR', + help='learning rate (default: 1.0)') + parser.add_argument('--gamma', type=float, default=0.7, metavar='M', + help='Learning rate step gamma (default: 0.7)') + parser.add_argument('--no-cuda', action='store_true', default=False, + help='disables CUDA training') + parser.add_argument('--dry-run', action='store_true', default=False, + help='quickly check a single pass') + parser.add_argument('--seed', type=int, default=1, metavar='S', + help='random seed (default: 1)') + parser.add_argument('--log-interval', type=int, default=10, metavar='N', + help='how many batches to wait before logging training status') + parser.add_argument('--sparsity', type=float, default=0.5, + help='target overall target sparsity') + args = parser.parse_args() + + main(args) diff --git a/examples/model_compress/pruning/multi_prune_torch.py b/examples/model_compress/multi_compressor_torch.py similarity index 97% rename from examples/model_compress/pruning/multi_prune_torch.py rename to examples/model_compress/multi_compressor_torch.py index 0cc2e36982..c24e9f9ebc 100644 --- a/examples/model_compress/pruning/multi_prune_torch.py +++ b/examples/model_compress/multi_compressor_torch.py @@ -1,11 +1,6 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -''' -NNI example for quick start of pruning. -In this example, we use level pruner to prune the LeNet on MNIST. -''' - import logging import argparse @@ -15,7 +10,7 @@ import torch.optim as optim from torchvision import datasets, transforms from torch.optim.lr_scheduler import StepLR -from models.mnist.lenet import LeNet +from pruning.models.mnist.lenet import LeNet from nni.algorithms.compression.pytorch.multicompressor import MultiCompressor import nni @@ -104,7 +99,7 @@ def main(args): 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) @@ -170,7 +165,7 @@ def main(args): model = pruner.compress() if __name__ == '__main__': - # Training settings + # Training settings parser = argparse.ArgumentParser(description='PyTorch MNIST Example for model comporession') parser.add_argument('--batch-size', type=int, default=64, metavar='N', help='input batch size for training (default: 64)') diff --git a/nni/algorithms/compression/pytorch/autocompressor.py b/nni/algorithms/compression/pytorch/autocompressor.py new file mode 100644 index 0000000000..27532742be --- /dev/null +++ b/nni/algorithms/compression/pytorch/autocompressor.py @@ -0,0 +1,108 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import copy +import logging +import os +import time +import torch + +from collections import OrderedDict +from pathlib import Path +from typing import Tuple, Union, List, Optional + +from nni.algorithms.compression.pytorch.multicompressor import MultiCompressor + +_logger = logging.getLogger(__name__) + + +LAYER_PRUNER_MAPPING = { + 'Conv2d': ['l1', 'l2'], + 'Linear': ['level'], + 'BatchNorm2d': ['slim'] +} + +LAYER_QUANTIZER_MAPPING = { + 'Conv2d': ['qat'], + 'Linear': ['qat'] +} + +class AutoCompressor: + def __init__(self, model, sparsity: float = 0.5, optimizer=None, trainer=None): + self.bound_model = model + self.sparsity = sparsity + self.optimizer = optimizer + self.trainer = trainer + self.pruner_layer_dict, self.quantizer_layer_dict, self.max_count = self._detect_search_space() + self.counter = 0 + + def _detect_search_space(self) -> Tuple[OrderedDict, OrderedDict, int]: + pruner_layer = set() + quantizer_layer = set() + for _, module in self.bound_model.named_modules(): + layer_type = type(module).__name__ + if module == self.bound_model: + continue + if layer_type in LAYER_PRUNER_MAPPING: + pruner_layer.add(layer_type) + else: + _logger.debug('Unsupported auto pruning layer: %s', layer_type) + if layer_type in LAYER_QUANTIZER_MAPPING: + quantizer_layer.add(layer_type) + else: + _logger.debug('Unsupported auto quantizing layer: %s', layer_type) + + assert len(pruner_layer) + len(quantizer_layer) > 0, 'The model has no supported layer to compress.' + + total_combination_num = 1 + + pruner_layer_dict = OrderedDict() + for layer_name in pruner_layer: + pruner_layer_dict.setdefault(layer_name, []) + for pruner_type in LAYER_PRUNER_MAPPING[layer_name]: + pruner_layer_dict[layer_name].append((pruner_type, {}, {'sparsity': self.sparsity, 'op_types': [layer_name]})) + total_combination_num *= len(pruner_layer_dict[layer_name]) + + quantizer_layer_dict = OrderedDict() + for layer_name in quantizer_layer: + quantizer_layer_dict.setdefault(layer_name, []) + for quantizer_type in LAYER_QUANTIZER_MAPPING[layer_name]: + quantizer_layer_dict[layer_name].append((quantizer_type, {}, {'quant_types': ['weight'], 'quant_bits': {'weight': 8}, 'op_types': [layer_name]})) + total_combination_num *= len(quantizer_layer_dict[layer_name]) + + return pruner_layer_dict, quantizer_layer_dict, int(total_combination_num) + + def _generate_config_list(self): + quo, rem = self.counter, 0 + pruner_config_dict = {} + for layer_name, choices in self.pruner_layer_dict.items(): + quo, rem = quo // len(choices), quo % len(choices) + name, args, config = choices[rem] + pruner_config_dict.setdefault(name, {'config_list': [], 'pruner': {'type': name, 'args': args}}) + pruner_config_dict[name]['config_list'].append(config) + + quantizer_config_dict = {} + for layer_name, choices in self.quantizer_layer_dict.items(): + quo, rem = quo // len(choices), quo % len(choices) + name, args, config = choices[rem] + quantizer_config_dict.setdefault(name, {'config_list': [], 'quantizer': {'type': name, 'args': args}}) + quantizer_config_dict[name]['config_list'].append(config) + + self.counter += 1 + + return list(pruner_config_dict.values()) + list(quantizer_config_dict.values()) + + def run(self, result_dir: str = None, input_shape: Optional[Union[List, Tuple]] = None, device: torch.device = None): + if result_dir is None: + result_dir = os.path.abspath('./autocompress_result_{}'.format(int(time.time()))) + Path(result_dir).mkdir(parents=True, exist_ok=True) + while self.counter < self.max_count: + mixed_config_list = self._generate_config_list() + config_result_dir = os.path.join(result_dir, self.counter) + _logger.info('Result saved under %s', config_result_dir) + + compressor = MultiCompressor(copy.deepcopy(self.bound_model), mixed_config_list, optimizer=copy.deepcopy(self.optimizer), trainer=self.trainer) + compressor.set_config(os.path.join(config_result_dir, 'model.pt'), os.path.join(config_result_dir, 'mask.pt'), + os.path.join(config_result_dir, 'calibration.pt'), os.path.join(config_result_dir, 'onnx.pt'), + input_shape=input_shape, device=device) + compressor.compress() diff --git a/nni/algorithms/compression/pytorch/multicompressor.py b/nni/algorithms/compression/pytorch/multicompressor.py index 7bb2f6b9d9..6c98c6b850 100644 --- a/nni/algorithms/compression/pytorch/multicompressor.py +++ b/nni/algorithms/compression/pytorch/multicompressor.py @@ -8,8 +8,8 @@ from nni.compression.pytorch.compressor import Compressor from nni.compression.pytorch.speedup import ModelSpeedup -from nni.algorithms.compression.pytorch.pruning.one_shot import LevelPruner, SlimPruner, L1FilterPruner, L2FilterPruner, FPGMPruner, - TaylorFOWeightFilterPruner, ActivationAPoZRankFilterPruner, ActivationMeanRankFilterPruner +from nni.algorithms.compression.pytorch.pruning.one_shot import LevelPruner, SlimPruner, L1FilterPruner, L2FilterPruner, \ + FPGMPruner, TaylorFOWeightFilterPruner, ActivationAPoZRankFilterPruner, ActivationMeanRankFilterPruner from nni.algorithms.compression.pytorch.quantization.quantizers import NaiveQuantizer, QAT_Quantizer, DoReFaQuantizer, BNNQuantizer _logger = logging.getLogger(__name__) @@ -60,7 +60,7 @@ def set_config(self, model_path: str, mask_path: str = None, calibration_path: s self.calibration_path = calibration_path self.onnx_path = onnx_path self.input_shape = input_shape - self.device = device + self.device = device if device else 'cpu' def _convert_config_list(self, mixed_config_list: List[dict]) -> Tuple[CONCFIG_LIST_TYPE, CONCFIG_LIST_TYPE]: pruner_config_list = [] From 569079ebae9b7328072063da171d3300c5892244 Mon Sep 17 00:00:00 2001 From: J-shang Date: Mon, 26 Apr 2021 11:03:31 +0800 Subject: [PATCH 10/10] temp sync --- nni/algorithms/compression/pytorch/autocompressor.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nni/algorithms/compression/pytorch/autocompressor.py b/nni/algorithms/compression/pytorch/autocompressor.py index 27532742be..d34ac4089f 100644 --- a/nni/algorithms/compression/pytorch/autocompressor.py +++ b/nni/algorithms/compression/pytorch/autocompressor.py @@ -95,10 +95,10 @@ def _generate_config_list(self): def run(self, result_dir: str = None, input_shape: Optional[Union[List, Tuple]] = None, device: torch.device = None): if result_dir is None: result_dir = os.path.abspath('./autocompress_result_{}'.format(int(time.time()))) - Path(result_dir).mkdir(parents=True, exist_ok=True) while self.counter < self.max_count: mixed_config_list = self._generate_config_list() - config_result_dir = os.path.join(result_dir, self.counter) + config_result_dir = os.path.join(result_dir, str(self.counter)) + Path(config_result_dir).mkdir(parents=True, exist_ok=True) _logger.info('Result saved under %s', config_result_dir) compressor = MultiCompressor(copy.deepcopy(self.bound_model), mixed_config_list, optimizer=copy.deepcopy(self.optimizer), trainer=self.trainer)