import keras fix + install on AWS DL AMI fix (#10)

* Fixed bug whereby pruning would not be done if keras was imported via `import keras`

* Added additional unit tests to ensure that LTP works when keras is imported in various different ways.

* Changed setup.py to allow this package to be installed in AWS deep learning AMIs where tensorflow" package is named "tensorflow-gpu"

Co-authored-by: jim.meyer <no-reply>

.flake8

@@ -1,3 +1,4 @@
 [flake8]
 max-line-length = 120
 ignore = E128
+exclude = .git,.tox,__pycache__,build,dist

.github/workflows/std-build-test.yml

@@ -27,7 +27,8 @@ jobs:
         flake8 . --count --show-source
     - name: Run unit tests
       run: |
-        python -m pytest --cov=lottery_ticket_pruner --cov-branch --junitxml build/unittest_results-${{ matrix.python-version }}.xml --cov-report html:build/test_cov-${{ matrix.python-version }}
+        python -m pytest --cov=lottery_ticket_pruner --cov-branch --junitxml build/unittest_results-${{ matrix.python-version }}.xml --ignore-glob 'tests/*_randseed.py'
+        python -m pytest --cov=lottery_ticket_pruner --cov-branch --junitxml build/unittest_results-${{ matrix.python-version }}-randseed.xml --cov-append --cov-report html:build/test_cov-${{ matrix.python-version }} tests/*_randseed.py
     - name: Upload test results
       uses: actions/upload-artifact@v1
       with:
@@ -62,7 +63,8 @@ jobs:
         flake8 . --count --show-source
     - name: Run unit tests
       run: |
-        python -m pytest --cov=lottery_ticket_pruner --cov-branch --junitxml build/unittest_results-${{ matrix.python-version }}.xml --cov-report html:build/test_cov-${{ matrix.python-version }}
+        python -m pytest --cov=lottery_ticket_pruner --cov-branch --junitxml build/unittest_results-${{ matrix.python-version }}.xml --ignore-glob 'tests/*_randseed.py'
+        python -m pytest --cov=lottery_ticket_pruner --cov-branch --junitxml build/unittest_results-${{ matrix.python-version }}-randseed.xml --cov-append --cov-report html:build/test_cov-${{ matrix.python-version }} tests/*_randseed.py
     - name: Upload test results
       uses: actions/upload-artifact@v1
       with:

RELEASE.md

@@ -1,4 +1,13 @@
-# 0.1.0
+# 0.8.1
+
+Changed setup.py to allow this package to be installed in AWS deep learning AMIs where "tensorflow" package is named
+"tensorflow-gpu".
+
+Fixed bug whereby pruning would not be done if keras was imported via `import keras`. Using
+`import tensorflow.keras as keras` or `from tensorflow.python import keras` were working fine though.
+Added unit tests that import keras in various different ways to ensure this package works regardless.
+
+# 0.8.0
 
 Initial functional package. Tested via unit tests and via integration into an unrelated image classification pipeline
 based on MobilenetV2.

lottery_ticket_pruner/__init__.py

@@ -1,4 +1,4 @@
 from .lottery_ticket_pruner import LotteryTicketPruner  # noqa
 from .keras_pruner_callback import PrunerCallback  # noqa
 
-__version__ = '0.8.0'
+__version__ = '0.8.1'

lottery_ticket_pruner/lottery_ticket_pruner.py

@@ -2,7 +2,6 @@
 import math
 import sys
 
-import tensorflow.keras as keras
 import numpy as np
 
 logger = logging.getLogger('lottery_ticket_pruner')
@@ -224,6 +223,22 @@ def __init__(self, initial_model):
             which weights of the model being trained are to be pruned.
             Some pruning strategies require access to the initial weights of the model to determine the pruning mask.
         """
+        self.prunable_layer_names = ('Conv1D',
+                                     'Conv2D',
+                                     'Conv2DTranspose',
+                                     'Conv3D',
+                                     'Conv3DTranspose',
+                                     'Convolution1D',
+                                     'Convolution2D',
+                                     'Convolution2DTranspose',
+                                     'Convolution3D',
+                                     'Convolution3DTranspose',
+                                     'Dense',
+                                     'DepthwiseConv2D',
+                                     'SeparableConv1D',
+                                     'SeparableConv2D',
+                                     'SeparableConvolution1D',
+                                     'SeparableConvolution2D')
 
         # Now determine which weights of which layers are prunable
         layer_index = 0
@@ -279,18 +294,17 @@ def _verify_compatible_model(self, model):
                         layer.name, getattr(layer, 'output_shape', None), output_shape))
 
     def _prunable(self, layer, weights):
-        return isinstance(layer, (keras.layers.Conv1D,
-                                  keras.layers.SeparableConv1D,
-                                  keras.layers.Conv2D,
-                                  keras.layers.Conv2DTranspose,
-                                  # keras.layers.Convolution2DTranspose,    an alias for keras.layers.Conv2DTranspose
-                                  keras.layers.Convolution2D,
-                                  keras.layers.DepthwiseConv2D,
-                                  keras.layers.SeparableConv2D,
-                                  keras.layers.Conv3D,
-                                  keras.layers.Convolution3D,
-                                  keras.layers.Dense,
-                                  )) and len(weights.shape) > 1
+        """ Depending on how callers import keras and what version of tensorflow is being used  the package path to the
+         layers we're interested in can be any of (at least):
+                keras.layers.Conv1D
+                keras.layers.convolutational.Conv1D
+                tensorflow.python.keras.layers.convolutational.Conv1D
+         Hence we use some relaxed logic here.
+         This is unfortunate since it potentially excludes subclasses of any of the supported keras layers from being
+         prunable.
+         """
+        return len(weights.shape) > 1 and 'keras.layers' in type(layer).__module__ and type(
+            layer).__name__ in self.prunable_layer_names
 
     def apply_dwr(self, model):
         """ Applies Dynamic Weight Rescaling (DWR) to the unpruned weights in the model.

setup.py

@@ -28,7 +28,10 @@ def get_long_description():
     long_description_content_type='text/markdown',
     url='https://github.com/jim-meyer/lottery_ticket_pruner',
     packages=setuptools.find_packages(),
-    install_requires=['keras>=2.1.0', 'tensorflow>=1.12', 'numpy>=1.18.3'],
+    # Don't specify version of tensorflow so that this package can easily be used in AWS deep learning AMIs
+    # where the "tensorflow" package is actually named "tensorflow-gpu"
+    # install_requires=['keras>=2.1.0', 'tensorflow>=1.12', 'numpy>=1.18.3'],
+    install_requires=['keras>=2.1.0', 'numpy>=1.18.3'],
     classifiers=[
         'Programming Language :: Python :: 3',
         "Development Status :: 4 - Beta",

tests/test_from_tensorflow_python_import_keras.py

@@ -0,0 +1,30 @@
+import unittest
+
+# We just make sure that we're finding all the expected prunable layers since layer names are different
+# depending on how we import keras. Some of the different ways that keras can be imported:
+#   import tensorflow.keras as keras
+#   from tensorflow.python import keras
+#   import keras
+from tensorflow.python import keras
+
+import lottery_ticket_pruner
+
+
+class TestFromTensorflowPythonImportKeras(unittest.TestCase):
+    def test_inception_v3(self):
+        if hasattr(keras.applications, 'InceptionV3'):
+            factory_func = keras.applications.InceptionV3
+        elif hasattr(keras.applications.inception_v3, 'InceptionV3'):
+            factory_func = keras.applications.inception_v3.InceptionV3
+        else:
+            raise Exception('Cannot find InceptionV3 while using `from tensorflow.python import keras`')
+        model = factory_func(input_shape=(299, 299, 3),
+                             weights='imagenet',
+                             include_top=True,
+                             pooling='max')
+        pruner = lottery_ticket_pruner.LotteryTicketPruner(model)
+        self.assertEqual(95, len(pruner.prune_masks_map))
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_import_keras.py

@@ -0,0 +1,29 @@
+import unittest
+
+# We just make sure that we're finding all the expected prunable layers since layer names are different
+# depending on how we import keras. Some of the different ways that keras can be imported:
+#   import tensorflow.keras as keras
+#   from tensorflow.python import keras
+#   import keras
+try:
+    import keras
+    KERAS_IMPORTED = True
+except ImportError:
+    KERAS_IMPORTED = False
+
+import lottery_ticket_pruner
+
+
+@unittest.skipIf(not KERAS_IMPORTED, 'Skipping unit tests that uses `import keras` since keras is not installed per se')
+class TestImportKeras(unittest.TestCase):
+    def test_inception_v3(self):
+        model = keras.applications.InceptionV3(input_shape=(299, 299, 3),
+                                               weights='imagenet',
+                                               include_top=True,
+                                               pooling='max')
+        pruner = lottery_ticket_pruner.LotteryTicketPruner(model)
+        self.assertEqual(95, len(pruner.prune_masks_map))
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_import_tensorflow_keras_as_keras.py

@@ -0,0 +1,24 @@
+import unittest
+
+# We just make sure that we're finding all the expected prunable layers since layer names are different
+# depending on how we import keras. Some of the different ways that keras can be imported:
+#   import tensorflow.keras as keras
+#   from tensorflow.python import keras
+#   import keras
+import tensorflow.keras as keras
+
+import lottery_ticket_pruner
+
+
+class TestImportTensorflowKerasAsKeras(unittest.TestCase):
+    def test_mobilenet_v2_from_tf2(self):
+        model = keras.applications.MobileNetV2(input_shape=(224, 224, 3),
+                                               weights='imagenet',
+                                               include_top=True,
+                                               pooling='max')
+        pruner = lottery_ticket_pruner.LotteryTicketPruner(model)
+        self.assertEqual(53, len(pruner.prune_masks_map))
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_keras_pruner_callback.py

@@ -9,7 +9,7 @@
 TEST_PRUNE_RATE = 0.5
 
 
-class TestVerificationCallback(keras.callbacks.Callback):
+class _VerificationCallback(keras.callbacks.Callback):
     """ Does verifications for these tests """
     def __init__(self, testcase):
         super().__init__()
@@ -69,7 +69,7 @@ def _assert_weights_have_been_pruned(self):
     def test_callback(self):
         epochs = 2
         # Can't do this in constructor of MNISTTest() since we have a chicken and egg problem
-        self.mnist_test.init(self.pruner, TestVerificationCallback(self))
+        self.mnist_test.init(self.pruner, _VerificationCallback(self))
         self.pruner.calc_prune_mask(self.model, TEST_PRUNE_RATE, 'smallest_weights')
 
         self.mnist_test.fit(self.model, epochs)

tests/test_lottery_ticker_pruner_randseed.py

@@ -0,0 +1,141 @@
+# This test does specific comparisons that depend on the random seeds being set the same from run to run
+import random
+random.seed(1234)
+
+import numpy as np  # noqa
+np.random.seed(2345)
+
+# Dancing needed to work with TF 1.x and 2.x
+import tensorflow   # noqa
+if hasattr(tensorflow, 'set_random_seed'):
+    tensorflow.set_random_seed(3456)
+else:
+    tensorflow.random.set_seed(3456)
+
+import unittest     # noqa
+
+import numpy as np  # noqa
+import tensorflow.keras as keras    # noqa
+
+import lottery_ticket_pruner    # noqa
+
+TEST_DNN_INPUT_DIMS = (64, 64, 3)
+TEST_DNN_NUM_CLASSES = 10
+
+
+class TestLotteryTicketPrunerRandseed(unittest.TestCase):
+    def _create_test_dnn_model(self):
+        input = keras.Input(shape=TEST_DNN_INPUT_DIMS, dtype='float32')
+        x = keras.layers.Conv2D(4,
+                                kernel_size=3,
+                                strides=(2, 2),
+                                padding='valid',
+                                use_bias=True,
+                                name='Conv1')(input)
+        x = keras.layers.BatchNormalization(axis=1,
+                                            epsilon=1e-3,
+                                            momentum=0.999,
+                                            name='bn_Conv1')(x)
+        x = keras.layers.ReLU(6., name='Conv1_relu')(x)
+
+        x = keras.layers.Conv2D(3,
+                                kernel_size=1,
+                                padding='same',
+                                use_bias=False,
+                                activation=None,
+                                name='Conv2')(x)
+        x = keras.layers.BatchNormalization(axis=1,
+                                      epsilon=1e-3,
+                                      momentum=0.999,
+                                      name='bn_Conv2')(x)
+        x = keras.layers.ReLU(6., name='Conv2_relu')(x)
+
+        x = keras.layers.GlobalAveragePooling2D()(x)
+        x = keras.layers.Dense(TEST_DNN_NUM_CLASSES, activation='softmax',
+                               use_bias=True, name='Logits')(x)
+        model = keras.Model(inputs=input, outputs=x)
+        return model
+
+    #
+    # calc_prune_mask()
+    #   'smallest_weights_global'
+    #
+    def test_smallest_weights_global(self):
+        """ Tests case where many or all weights are same value. Hence we might be tempted to mask on all of the
+        smallest weights rather than honoring only up to the prune rate
+        """
+        model = self._create_test_dnn_model()
+        interesting_layers = [model.layers[1], model.layers[4], model.layers[8]]
+        interesting_weights_index = 0
+
+        # Make sure no weights are zero so our checks below for zeroes only existing in masked weights are reliable
+        weight_counts = []
+        for layer in interesting_layers:
+            weights = layer.get_weights()
+            weights[interesting_weights_index][weights[interesting_weights_index] == 0.0] = 0.1234
+            layer.set_weights(weights)
+            num_weights = np.prod(weights[interesting_weights_index].shape)
+            weight_counts.append(num_weights)
+
+        pruner = lottery_ticket_pruner.LotteryTicketPruner(model)
+
+        num_pruned1 = 0
+        for layer in interesting_layers:
+            weights = layer.get_weights()
+            num_pruned1 += np.sum(weights[interesting_weights_index] == 0.0)
+
+        prune_rate = 0.5
+        pruner.calc_prune_mask(model, prune_rate, 'smallest_weights_global')
+
+        # calc_prune_mask() shouldn't do the actual pruning so verify that weights didn't change
+        num_pruned2 = 0
+        for layer in interesting_layers:
+            weights = layer.get_weights()
+            num_pruned2 += np.sum(weights[interesting_weights_index] == 0.0)
+        self.assertEqual(num_pruned1, num_pruned2)
+
+        pruner.apply_pruning(model)
+        pruned_counts = []
+        for layer in interesting_layers:
+            weights = layer.get_weights()
+            pruned_counts.append(np.sum(weights[interesting_weights_index] == 0.0))
+
+        total_weights = np.sum(weight_counts)
+        num_pruned = np.sum(pruned_counts)
+        self.assertAlmostEqual(prune_rate, num_pruned / total_weights, places=1)
+        # Given the seeding we did at the beginning of this test these results should be reproducible. They were
+        # obtained by manual inspection.
+        # Ranges are used here since TF 1.x on python 3.6, 3.7 gives slightly different results from TF 2.x on
+        # python 3.8. These assertions accomodate both.
+        self.assertTrue(62 <= pruned_counts[0] <= 67, msg=f'pruned_counts={pruned_counts}')
+        self.assertTrue(2 <= pruned_counts[1] <= 5, msg=f'pruned_counts={pruned_counts}')
+        self.assertTrue(5 <= pruned_counts[2] <= 9, msg=f'pruned_counts={pruned_counts}')
+        self.assertEqual(75, sum(pruned_counts))
+
+        # Now prune once more to make sure cumulative pruning works as expected
+        total_prune_rate = prune_rate
+        prune_rate = 0.2
+        total_prune_rate = total_prune_rate + (1.0 - total_prune_rate) * prune_rate
+        pruner.calc_prune_mask(model, prune_rate, 'smallest_weights_global')
+        pruner.apply_pruning(model)
+
+        pruned_counts = []
+        for layer in interesting_layers:
+            weights = layer.get_weights()
+            pruned_counts.append(np.sum(weights[interesting_weights_index] == 0.0))
+
+        total_weights = np.sum(weight_counts)
+        num_pruned = np.sum(pruned_counts)
+        self.assertEqual(num_pruned / total_weights, total_prune_rate)
+        # Given the seeding we did at the beginning of this test these results should be reproducible. They were
+        # obtained by manual inspection.
+        # Ranges are used here since TF 1.x on python 3.6, 3.7 gives slightly different results from TF 2.x on
+        # python 3.8. These assertions accomodate both.
+        self.assertTrue(74 <= pruned_counts[0] <= 78, msg=f'pruned_counts={pruned_counts}')
+        self.assertTrue(2 <= pruned_counts[1] <= 5, msg=f'pruned_counts={pruned_counts}')
+        self.assertTrue(9 <= pruned_counts[2] <= 12, msg=f'pruned_counts={pruned_counts}')
+        self.assertEqual(90, sum(pruned_counts))
+
+
+if __name__ == '__main__':
+    unittest.main()