TF compression fix and UT (#2817)

* draft

* fix bug

* add ut

* fix typo

* skip tfv1 ut environments
Co-authored-by: liuzhe <zhe.liu@microsoft.com>

examples/model_compress/model_prune_tf.py

@@ -28,7 +28,12 @@ def get_dataset(dataset_name='mnist'):
 
 def create_model(model_name='naive'):
     assert model_name == 'naive'
-    return tf.keras.Sequential([
+    return NaiveModel()
+
+class NaiveModel(tf.keras.Model):
+    def __init__(self):
+        super().__init__()
+        self.seq_layers = [
             tf.keras.layers.Conv2D(filters=20, kernel_size=5),
             tf.keras.layers.BatchNormalization(),
             tf.keras.layers.ReLU(),
@@ -42,7 +47,12 @@ def create_model(model_name='naive'):
             tf.keras.layers.ReLU(),
             tf.keras.layers.Dense(units=10),
             tf.keras.layers.Softmax()
-    ])
+        ]
+
+    def call(self, x):
+        for layer in self.seq_layers:
+            x = layer(x)
+        return x
 
 
 def create_pruner(model, pruner_name):
@@ -57,18 +67,38 @@ def main(args):
     train_set, test_set = get_dataset(dataset_name)
     model = create_model(model_name)
 
-    optimizer = tf.keras.optimizers.SGD(learning_rate=0.1, momentum=0.9, decay=1e-4)
-    model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])
-
     print('start training')
-    model.fit(train_set[0], train_set[1], batch_size=args.batch_size, epochs=args.pretrain_epochs, validation_data=test_set)
+    optimizer = tf.keras.optimizers.SGD(learning_rate=0.1, momentum=0.9, decay=1e-4)
+    model.compile(
+        optimizer=optimizer,
+        loss='sparse_categorical_crossentropy',
+        metrics=['accuracy']
+    )
+    model.fit(
+        train_set[0],
+        train_set[1],
+        batch_size=args.batch_size,
+        epochs=args.pretrain_epochs,
+        validation_data=test_set
+    )
 
     print('start model pruning')
     optimizer_finetune = tf.keras.optimizers.SGD(learning_rate=0.001, momentum=0.9, decay=1e-4)
     pruner = create_pruner(model, args.pruner_name)
     model = pruner.compress()
-    model.compile(optimizer=optimizer_finetune, loss='sparse_categorical_crossentropy', metrics=['accuracy'])
-    model.fit(train_set[0], train_set[1], batch_size=args.batch_size, epochs=args.prune_epochs, validation_data=test_set)
+    model.compile(
+        optimizer=optimizer_finetune,
+        loss='sparse_categorical_crossentropy',
+        metrics=['accuracy'],
+        run_eagerly=True  # NOTE: Important, model compression does not work in graph mode!
+    )
+    model.fit(
+        train_set[0],
+        train_set[1],
+        batch_size=args.batch_size,
+        epochs=args.prune_epochs,
+        validation_data=test_set
+    )
 
 
 if __name__ == '__main__':

src/sdk/pynni/nni/compression/tensorflow/compressor.py

@@ -6,7 +6,10 @@ Abstract base classes for TensorFlow model compression.
 """
 
 import logging
+
 import tensorflow as tf
+assert tf.__version__.startswith('2'), 'NNI model compression only supports TensorFlow v2.x'
+
 from . import default_layers
 
 _logger = logging.getLogger(__name__)
@@ -25,9 +28,9 @@ class LayerInfo:
         The layer's name. Note that it's local to sub-model and may differ from its attribute name.
     type : str
         Name of the layer's class.
-    path : list of str/int
+    path : list of str or tuple of (str, int)
         The layer object's and its parents' attribute name / list index.
-        For example, if the path is `['cells', 2, 'conv']`, then the layer can be accessed as `model.cells[2].conv`.
+        For example, if the path is `[('cells', 2), 'conv']`, then the layer can be accessed as `model.cells[2].conv`.
     config : JSON object
         Selected configuration for this layer. The format is detailed in tutorial.
 
@@ -35,7 +38,7 @@ class LayerInfo:
     ----------
     layer : tf.keras.layers.Layer
         See attributes section.
-    path : list of str/int
+    path : list of str or tuple of (str, int)
         See attributes section.
     """
 
@@ -75,6 +78,8 @@ class Compressor:
 
     def __init__(self, LayerWrapperClass, model, config_list):
         assert isinstance(model, tf.keras.Model)
+        if isinstance(model, tf.keras.Sequential):
+            raise ValueError('NNI model compression does not support `Sequential` model for now')
         self.validate_config(model, config_list)
 
         self.bound_model = model
@@ -204,10 +209,12 @@ class PrunerLayerWrapper(tf.keras.Model):
         for weight in self.layer.weights:
             mask = self.masks.get(weight.name)
             if mask is not None:
-                new_weights.append(tf.math.multiply(weight, mask).numpy())
+                new_weights.append(tf.math.multiply(weight, mask))
             else:
-                new_weights.append(weight.numpy())
-        self.layer.set_weights(new_weights)
+                new_weights.append(weight)
+        if new_weights and not hasattr(new_weights[0], 'numpy'):
+            raise RuntimeError('NNI: Compressed model can only run in eager mode')
+        self.layer.set_weights([weight.numpy() for weight in new_weights])
         return self.layer(*inputs)
 
 
@@ -244,26 +251,21 @@ def _locate_layers(model, cur_path=[]):
     # and to my knowledge `Layer.name` is only useful for read-only access.
     # `cur_path`s format is documented in `LayerInfo.path`.
     # TODO: it can only find layers in `Model` and `list` for now.
+    assert isinstance(model, tf.keras.Model)
+    if isinstance(model, tf.keras.Sequential):
+        _logger.warning('`Sequential` model is not supported yet, ignored.')
     ret = {}
-
-    if isinstance(model, tf.keras.Model):
     for key, value in model.__dict__.items():
         if isinstance(value, tf.keras.Model):
             ret.update(_locate_layers(value, cur_path + [key]))
-            elif isinstance(value, list):
-                ret.update(_locate_layers(value, cur_path + [key]))
         elif isinstance(value, tf.keras.layers.Layer):
             ret[id(value)] = LayerInfo(value, cur_path + [key])
-
-    elif isinstance(model, list):
-        for i, item in enumerate(model):
+        elif isinstance(value, list):
+            for i, item in enumerate(value):
                 if isinstance(item, tf.keras.Model):
-                ret.update(_locate_layers(item, cur_path + [i]))
+                    ret.update(_locate_layers(item, cur_path + [(key, i)]))
                 elif isinstance(item, tf.keras.layers.Layer):
-                ret[id(item)] = LayerInfo(item, cur_path + [i])
-
-    else:
-        raise ValueError('Unexpected model type: {}'.format(type(model)))
+                    ret[id(item)] = LayerInfo(item, cur_path + [(key, i)])
     return ret
 
 def _select_config(layer_info, config_list):
@@ -289,12 +291,17 @@ def _instrument_model(model, wrappers):
     for wrapper in reversed(wrappers):
         cur = model
         for key in wrapper.layer_info.path[:-1]:
-            if isinstance(key, int):
-                cur = cur[key]
-            else:
+            if isinstance(key, str):
                 cur = getattr(cur, key)
-        key = wrapper.layer_info.path[-1]
-        if isinstance(key, int):
-            cur[key] = wrapper
             else:
+                name, index = key
+                cur = getattr(cur, name)[index]
+        key = wrapper.layer_info.path[-1]
+        if isinstance(key, str):
             setattr(cur, key, wrapper)
+        else:
+            name, index = key
+            getattr(cur, name)[index] = wrapper
+            #if isinstance(cur, tf.keras.Sequential):
+            #    cur._graph_initialized = False
+            #    cur._layer_call_argspecs[wrapper] = cur._layer_call_argspecs[wrapper.layer]

src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py

@@ -44,20 +44,24 @@ class LevelPrunerMasker(WeightMasker):
     def calc_masks(self, sparsity, wrapper, wrapper_idx=None):
         masks = {}
         for weight_variable in wrapper.layer.weights:
-            if weight_variable.name == 'bias':
+            if 'bias' in weight_variable.name:
                 continue
 
-            k = int(tf.size(weight_variable).numpy() * sparsity)
-            if k == 0:
+            num_prune = int(tf.size(weight_variable).numpy() * sparsity)
+            if num_prune == 0:
                 continue
 
             weight = weight_variable.read_value()
             if wrapper.masks.get(weight_variable.name) is not None:
                 weight = tf.math.multiply(weight, wrapper.masks[weight_variable.name])
 
-            w_abs = tf.math.abs(tf.reshape(weight, [-1]))
-            threshold = tf.math.top_k(w_abs, k)[0][0]
-            mask = tf.math.greater(w_abs, threshold)
+            w_abs = tf.math.abs(weight)
+            k = tf.size(weight) - num_prune
+            topk = tf.math.top_k(tf.reshape(w_abs, [-1]), k)[0]
+            if tf.size(topk) == 0:
+                mask = tf.zeros_like(weight)
+            else:
+                mask = tf.math.greater_equal(w_abs, topk[-1])
             masks[weight_variable.name] = tf.cast(mask, weight.dtype)
         return masks
 

src/sdk/pynni/tests/test_compressor_tf.py

+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import unittest
+
+import numpy as np
+import tensorflow as tf
+
+
+####
+#
+# This file tests pruners on 2 models: a classic CNN model, and a naive model with one linear layer
+#
+# The CNN model is used to test layer detecting and instrumenting.
+#
+# The naive model is used to test mask calculation.
+# It has a single 10x10 linear layer without bias, and `reduce_sum` its result.
+# To help predicting pruning result, the linear layer has fixed initial weights:
+#     [ [ 0.0, 1.0, 2.0, ..., 9.0 ], [0.1, 1.1, 2.1, ..., 9.1 ], ... , [0.9, 1.0, 2.9, ..., 9.9 ] ]
+#
+####
+
+
+# This tensor is used as input of 10x10 linear layer, the first dimension is batch size
+tensor1x10 = tf.constant([[1.0] * 10])
+
+
+@unittest.skipIf(tf.__version__[0] != '2', 'Skip TF 1.x setup')
+class TfCompressorTestCase(unittest.TestCase):
+    def test_layer_detection(self):
+        # Conv and dense layers should be compressed, pool and flatten should not.
+        # This also tests instrumenting functionality.
+        self._test_layer_detection_on_model(CnnModel())
+
+    def _test_layer_detection_on_model(self, model):
+        pruner = pruners['level'](model)
+        pruner.compress()
+        layer_types = sorted(wrapper.layer_info.type for wrapper in pruner.wrappers)
+        assert layer_types == ['Conv2D', 'Dense', 'Dense'], layer_types
+
+    def test_level_pruner(self):
+        # prune 90% : 9.0 + 9.1 + ... + 9.9 = 94.5
+        model = build_naive_model()
+        pruners['level'](model).compress()
+        x = model(tensor1x10)
+        assert x.numpy() == 94.5
+
+
+try:
+    from tensorflow.keras import Model, Sequential
+    from tensorflow.keras.layers import (Conv2D, Dense, Flatten, MaxPool2D)
+
+    from nni.compression.tensorflow import LevelPruner
+
+    pruners = {
+        'level': (lambda model: LevelPruner(model, [{'sparsity': 0.9, 'op_types': ['default']}])),
+    }
+
+    class CnnModel(Model):
+        def __init__(self):
+            super().__init__()
+            self.conv = Conv2D(filters=10, kernel_size=3, activation='relu')
+            self.pool = MaxPool2D(pool_size=2)
+            self.flatten = Flatten()
+            self.fc1 = Dense(units=10, activation='relu')
+            self.fc2 = Dense(units=5, activation='softmax')
+
+        def call(self, x):
+            x = self.conv(x)
+            x = self.pool(x)
+            x = self.flatten(x)
+            x = self.fc1(x)
+            x = self.fc2(x)
+            return x
+
+    class NaiveModel(Model):
+        def __init__(self):
+            super().__init__()
+            self.fc = Dense(units=10, use_bias=False)
+
+        def call(self, x):
+            return tf.math.reduce_sum(self.fc(x))
+
+except Exception:
+    pass
+
+
+def build_naive_model():
+    model = NaiveModel()
+    model.build(tensor1x10.shape)
+    weight = [[(i + j * 0.1) for i in range(10)] for j in range(10)]
+    model.set_weights([np.array(weight)])
+    return model
+
+
+if __name__ == '__main__':
+    unittest.main()