From a40f4dc76ce6f22b5153e19f080aaff9a671ca8a Mon Sep 17 00:00:00 2001 From: Artur Mostowski Date: Mon, 16 Aug 2021 23:01:42 +0200 Subject: [PATCH] pruning | sparse layerwise method --- pruning/registry.py | 3 +- pruning/sparse_layerwise.py | 48 ++++++++++++++++ pruning/test/test_sparse_layerwise.py | 82 +++++++++++++++++++++++++++ 3 files changed, 132 insertions(+), 1 deletion(-) create mode 100644 pruning/sparse_layerwise.py create mode 100644 pruning/test/test_sparse_layerwise.py diff --git a/pruning/registry.py b/pruning/registry.py index ed284308..624be662 100644 --- a/pruning/registry.py +++ b/pruning/registry.py @@ -8,8 +8,9 @@ from foundations.hparams import PruningHparams from pruning import sparse_global +from pruning import sparse_layerwise -registered_strategies = {'sparse_global': sparse_global.Strategy} +registered_strategies = {'sparse_global': sparse_global.Strategy, 'sparse_layerwise': sparse_layerwise.Strategy} def get(pruning_hparams: PruningHparams): diff --git a/pruning/sparse_layerwise.py b/pruning/sparse_layerwise.py new file mode 100644 index 00000000..22c1397e --- /dev/null +++ b/pruning/sparse_layerwise.py @@ -0,0 +1,48 @@ +# Copyright (c) Facebook, Inc. and its affiliates. + +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import dataclasses +import numpy as np + +from foundations import hparams +import models.base +from pruning import base +from pruning import sparse_global +from pruning.mask import Mask + + +PruningHparams = sparse_global.PruningHparams + + +class Strategy(base.Strategy): + @staticmethod + def get_pruning_hparams() -> type: + return PruningHparams + + @staticmethod + def prune(pruning_hparams: PruningHparams, trained_model: models.base.Model, current_mask: Mask = None): + current_mask = Mask.ones_like(trained_model).numpy() if current_mask is None else current_mask.numpy() + + new_mask = Mask(current_mask) + + for k,v in current_mask.items(): + + if pruning_hparams.pruning_layers_to_ignore and k in pruning_hparams.pruning_layers_to_ignore: + continue + + # Determine the number of weights that need to be pruned. + number_of_remaining_weights = np.sum(v) + number_of_weights_to_prune = np.ceil( + pruning_hparams.pruning_fraction * number_of_remaining_weights).astype(int) + + + weights = trained_model.state_dict()[k].clone().cpu().detach().numpy() + + weight_vector = weights[current_mask[k] == 1] + threshold = np.sort(np.abs(weight_vector))[number_of_weights_to_prune] + + new_mask[k] = np.where(np.abs(weights) >= threshold, current_mask[k], np.zeros_like(v)) + + return new_mask diff --git a/pruning/test/test_sparse_layerwise.py b/pruning/test/test_sparse_layerwise.py new file mode 100644 index 00000000..59990084 --- /dev/null +++ b/pruning/test/test_sparse_layerwise.py @@ -0,0 +1,82 @@ +# Copyright (c) Facebook, Inc. and its affiliates. + +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np + +import models.registry +from pruning.sparse_layerwise import Strategy +from pruning.sparse_layerwise import PruningHparams +from testing import test_case + + +class TestSparseLayerWise(test_case.TestCase): + def setUp(self): + super(TestSparseLayerWise, self).setUp() + self.hparams = PruningHparams('sparse_global', 0.2) + + model_hparams = models.registry.get_default_hparams('cifar_resnet_20').model_hparams + self.model = models.registry.get(model_hparams) + + def test_get_pruning_hparams(self): + self.assertTrue(issubclass(Strategy.get_pruning_hparams(), PruningHparams)) + + def test_prune(self): + m = Strategy.prune(self.hparams, self.model) + + # Check that the mask only contains entries for the prunable layers. + self.assertEqual(set(m.keys()), set(self.model.prunable_layer_names)) + + # Check that the masks are the same sizes as the tensors. + for k in self.model.prunable_layer_names: + self.assertEqual(list(m[k].shape), list(self.model.state_dict()[k].shape)) + + # Check that the right fraction of weights was pruned for all prunable layers. + m = m.numpy() + for k, v in m.items(): + pruned_count = np.sum(1 - v) + should_be_pruned = v.size * self.hparams.pruning_fraction + self.assertTrue(should_be_pruned <= pruned_count <= should_be_pruned + 1) + + # Ensure that there are no errors in thresholding for every layer + for k, v in m.items(): + pruned_weights = self.model.state_dict()[k].numpy()[m[k] == 0] + threshold = np.max(np.abs(pruned_weights)) + unpruned_weights = self.model.state_dict()[k].numpy()[m[k] == 1] + self.assertTrue(np.all(np.abs(unpruned_weights) > threshold)) + + def test_iterative_pruning(self): + m = Strategy.prune(self.hparams, self.model) + m2 = Strategy.prune(self.hparams, self.model, m) + + # Ensure that all weights pruned before are still pruned here. + m, m2 = m.numpy(), m2.numpy() + self.assertEqual(set(m.keys()), set(m2.keys())) + for k in m: + self.assertTrue(np.all(m[k] >= m2[k])) + + for k, v in m2.items(): + pruned_count = np.sum(1 - v) + should_be_pruned = v.size * (1 - ((1 - self.hparams.pruning_fraction) ** 2)) + self.assertTrue(should_be_pruned <= pruned_count <= should_be_pruned + 2) + + def test_prune_layers_to_ignore(self): + layers_to_ignore = sorted(self.model.prunable_layer_names)[:5] + self.hparams.pruning_layers_to_ignore = ','.join(layers_to_ignore) + + m = Strategy.prune(self.hparams, self.model).numpy() + + # Ensure that the ignored layers were, indeed, ignored. + for k in layers_to_ignore: + self.assertTrue(np.all(m[k] == 1)) + + # Ensure that the expected fraction of parameters was still pruned. + total_pruned = np.sum([np.sum(1 - v) for v in m.values()]) + total_weights = np.sum([v.size for v in m.values()]) + actual_fraction = float(total_pruned) / total_weights + self.assertGreaterEqual(actual_fraction, self.hparams.pruning_fraction) + self.assertGreater(self.hparams.pruning_fraction + 0.0001, actual_fraction) + + +test_case.main()