From 21b25b7c07cd89139d007494a9f94a6ff414c344 Mon Sep 17 00:00:00 2001
From: liuzhe <zhliu1@microsoft.com>
Date: Fri, 31 Jul 2020 16:44:12 +0800
Subject: [PATCH 1/9] draft

---
 .../nni/compression/tensorflow/compressor.py  | 92 +++++++++++++------
 .../tensorflow/pruning/__init__.py            |  1 +
 .../tensorflow/pruning/one_shot.py            | 53 +++++++++++
 3 files changed, 120 insertions(+), 26 deletions(-)
 create mode 100644 src/sdk/pynni/nni/compression/tensorflow/pruning/__init__.py
 create mode 100644 src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py

diff --git a/src/sdk/pynni/nni/compression/tensorflow/compressor.py b/src/sdk/pynni/nni/compression/tensorflow/compressor.py
index 62580738a3..542f310f82 100644
--- a/src/sdk/pynni/nni/compression/tensorflow/compressor.py
+++ b/src/sdk/pynni/nni/compression/tensorflow/compressor.py
@@ -19,12 +19,18 @@ def __init__(self, keras_layer):
             self.weight = keras_layer.weights[self.weight_index]
         self._call = None
 
+def _setattr(model, name, module):
+    name_list = name.split('.')
+    for name in name_list[:-1]:
+        model = getattr(model, name)
+    setattr(model, name_list[-1], module)
+
 class Compressor:
     """
     Abstract base TensorFlow compressor
     """
 
-    def __init__(self, model, config_list):
+    def __init__(self, model, config_list, optimizer=None):
         """
         Record necessary info in class members
 
@@ -35,11 +41,29 @@ def __init__(self, model, config_list):
         config_list : list
             the configurations that users specify for compression
         """
+        assert isinstance(model, tf.keras.Model)
+        self.validate_config(model, config_list)
+
         self.bound_model = model
         self.config_list = config_list
-        self.modules_to_compress = []
+        self.optimizer = optimizer
+
+        self.modules_to_compress = None
+        self.modules_wrapper = []
+        self.is_wrapped = False
+
+        self._fwd_hook_handles = {}
+        self._fwd_hook_id = 0
+
+        for layer, config in self._detect_modules_to_compress():
+            wrapper = self._wrap_modules(layer, config)
+            self.modules_wrapper.append(wrapper)
+        if not self.modules_wrapper:
+            _logger.warning('Nothing is configured to compress, please check your model and config list')
 
-    def detect_modules_to_compress(self):
+        self._wrap_model()
+
+    def _detect_modules_to_compress(self):
         """
         detect all modules should be compressed, and save the result in `self.modules_to_compress`.
 
@@ -54,6 +78,11 @@ def detect_modules_to_compress(self):
                     self.modules_to_compress.append((layer, config))
         return self.modules_to_compress
 
+    def _wrap_model(self):
+        for wrapper in reversed(self.modules_wrapper):
+            _setattr(self.bound_model, wrapper.name, wrapper)
+        self.is_wrapped = True
+
     def compress(self):
         """
         Compress the model with algorithm implemented by subclass.
@@ -61,9 +90,9 @@ def compress(self):
         The model will be instrumented and user should never edit it after calling this method.
         `self.modules_to_compress` records all the to-be-compressed layers
         """
-        modules_to_compress = self.detect_modules_to_compress()
-        for layer, config in modules_to_compress:
-            self._instrument_layer(layer, config)
+        #modules_to_compress = self.detect_modules_to_compress()
+        #for layer, config in modules_to_compress:
+        #    self._instrument_layer(layer, config)
         return self.bound_model
 
     def get_modules_to_compress(self):
@@ -125,6 +154,9 @@ def step(self):
         """
 
 
+    def _wrap_modules(self, layer, config):
+        raise NotImplementedError()
+
     def _instrument_layer(self, layer, config):
         """
         This method is implemented in the subclasses, i.e., `Pruner` and `Quantizer`
@@ -156,7 +188,25 @@ class Pruner(Compressor):
     Abstract base TensorFlow pruner
     """
 
-    def calc_mask(self, layer, config):
+    def __init__(self, model, config_list, optimizer=None):
+        super().__init__(model, config_list, optimizer)
+        if optimizer is not None:
+            raise RuntimeError('Optimizer patching not implemented yet')
+            #self.patch_optimizer(self.update_mask)
+
+    def compress(self):
+        self.update_mask()
+        return self.bound_model
+
+    def update_mask(self):
+        for wrapper_idx, wrapper in enumerate(self.modules_wrapper):
+            masks = self.calc_mask(wrapper, wrapper_idx=wrapper_idx)
+            if masks is not None:
+                for k in masks:
+                    assert hasattr(wrapper, k)
+                    setattr(wrapper, k, masks[k])
+
+    def calc_mask(self, wrapper, **kwargs):
         """
         Pruners should overload this method to provide mask for weight tensors.
         The mask must have the same shape and type comparing to the weight.
@@ -172,7 +222,7 @@ def calc_mask(self, layer, config):
         """
         raise NotImplementedError("Pruners must overload calc_mask()")
 
-    def _instrument_layer(self, layer, config):
+    def _wrap_modules(self, layer, config):
         """
         Create a wrapper forward function to replace the original one.
 
@@ -183,22 +233,12 @@ def _instrument_layer(self, layer, config):
         config : dict
             the configuration for generating the mask
         """
-        layer._call = layer.keras_layer.call
-
-        def new_call(*inputs):
-            weights = [x.numpy() for x in layer.keras_layer.weights]
-            mask = self.calc_mask(layer, config)
-            weights[layer.weight_index] = weights[layer.weight_index] * mask
-            layer.keras_layer.set_weights(weights)
-            ret = layer._call(*inputs)
-            return ret
-
-        layer.keras_layer.call = new_call
-
-class Quantizer(Compressor):
-    """
-    Abstract base TensorFlow quantizer
-    """
+        return PrunerModuleWrapper(layer.module, layer.name, layer.type, config, self)
 
-    def quantize_weight(self, weight, config, op, op_type, op_name):
-        raise NotImplementedError("Quantizer must overload quantize_weight()")
+#class Quantizer(Compressor):
+#    """
+#    Abstract base TensorFlow quantizer
+#    """
+#
+#    def quantize_weight(self, weight, config, op, op_type, op_name):
+#        raise NotImplementedError("Quantizer must overload quantize_weight()")
diff --git a/src/sdk/pynni/nni/compression/tensorflow/pruning/__init__.py b/src/sdk/pynni/nni/compression/tensorflow/pruning/__init__.py
new file mode 100644
index 0000000000..f8ac8ea9b9
--- /dev/null
+++ b/src/sdk/pynni/nni/compression/tensorflow/pruning/__init__.py
@@ -0,0 +1 @@
+from .one_shot import *
diff --git a/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py b/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py
new file mode 100644
index 0000000000..0496b34983
--- /dev/null
+++ b/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py
@@ -0,0 +1,53 @@
+import tensorflow as tf
+
+from ..compressor import Pruner
+
+__all__ = [
+    'OneshotPruner'
+]
+
+class OneshotPruner(Pruner):
+    def __init__(self, model, config_list, pruning_algorithm='level', optimizer=None, **algo_kwargs):
+        super().__init__(model, config_list, optimizer)
+        self.set_wrapper_attribute('if_calculated', False)
+        self.masker = MASKER_DICT[pruning_alogrithm](model, self, **algo_kwargs)
+
+    def validate_config(self, model, config_list):
+        pass
+
+    def calc_mask(self, wrapper, wrapper_idx=None):
+        if wrapper.if_calculated:
+            return None
+
+        sparsity = wrapper.config['sparsity']
+        if not wrapper.if_calculated:
+            masks = self.masker.calc_mask(sparsity=sparsity, wrapper=wrapper, wrapper_idx=wrapper_idx)
+
+            if masks is not None:
+                wrapper.if_calculated = True
+            return masks
+        else:
+            return None
+
+
+MASKER_DICT = {
+    'level': LevelPrunerMasker,
+}
+
+class WeightMasker:
+    def __init__(self, model, pruner, **kwargs):
+        self.model = model
+        self.pruner = pruner
+
+    def calc_mask(self, sparsity, wrapper, wrapper_idx=None):
+        raise NotImplementedError()
+
+class LevelPrunerMasker(WeightMasker):
+    def calc_mask(self, sparsity, wrapper, wrapper_idx=None):
+        weight = wrapper.module.weight * wrapper.weight_mask
+        w_abs = tf.abs(w_abs)
+        k = int(tf.size(weight) * sparsity)
+        assert k > 0  # FIXME
+        threshold = tf.reduce_max(tf.topk(tf.reshape(w_abs, [-1]), k, largest=False)[0])
+        mask_weight = tf.cast((w_abs > threshold), weight.dtype)
+        return {'weight_mask': mask_weight}

From a5c82fdf01f93c92b0a36d7760e3f91b9032fe9a Mon Sep 17 00:00:00 2001
From: liuzhe <zhliu1@microsoft.com>
Date: Wed, 5 Aug 2020 10:38:55 +0800
Subject: [PATCH 2/9] update wrapper

---
 examples/model_compress/model_prune_tf.py     |  80 +++++++
 .../nni/compression/tensorflow/__init__.py    |   5 +-
 .../nni/compression/tensorflow/compressor.py  | 214 +++++-------------
 .../compression/tensorflow/default_layers.py  |  34 +--
 .../tensorflow/pruning/one_shot.py            |  27 ++-
 5 files changed, 161 insertions(+), 199 deletions(-)
 create mode 100644 examples/model_compress/model_prune_tf.py

diff --git a/examples/model_compress/model_prune_tf.py b/examples/model_compress/model_prune_tf.py
new file mode 100644
index 0000000000..3a928bc6da
--- /dev/null
+++ b/examples/model_compress/model_prune_tf.py
@@ -0,0 +1,80 @@
+import argparse
+
+import tensorflow as tf
+
+import nni.compression.tensorflow
+
+prune_config = {
+    'level': {
+        'dataset_name': 'mnist',
+        'model_name': 'naive',
+        'pruner_class': nni.compression.tensorflow.LevelPruner,
+        'config_list': [{
+            'sparsity': 0.5,
+            'op_types': ['default'],
+        }]
+    },
+}
+
+
+def get_dataset(dataset_name='mnist'):
+    assert dataset_name == 'mnist'
+
+    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
+    x_train = x_train[..., tf.newaxis] / 255.0
+    x_test = x_test[..., tf.newaxis] / 255.0
+    return (x_train, y_train), (x_test, y_test)
+
+
+def create_model(model_name='naive'):
+    assert model_name == 'naive'
+    return tf.keras.Sequential([
+        tf.keras.layers.Conv2D(filters=20, kernel_size=5),
+        tf.keras.layers.BatchNormalization(),
+        tf.keras.layers.ReLU(),
+        tf.keras.layers.MaxPool2D(pool_size=2),
+        tf.keras.layers.Conv2D(filters=20, kernel_size=5),
+        tf.keras.layers.BatchNormalization(),
+        tf.keras.layers.ReLU(),
+        tf.keras.layers.MaxPool2D(pool_size=2),
+        tf.keras.layers.Flatten(),
+        tf.keras.layers.Dense(units=500),
+        tf.keras.layers.ReLU(),
+        tf.keras.layers.Dense(units=10),
+    ])
+
+
+def create_pruner(model, pruner_name, optimizer=None):
+    pruner_class = prune_config[pruner_name]['pruner_class']
+    config_list = prune_config[pruner_name]['config_list']
+    return pruner_class(model, config_list, optimizer)
+
+
+def main(args):
+    model_name = prune_config[args.pruner_name]['model_name']
+    dataset_name = prune_config[args.pruner_name]['dataset_name']
+    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)
+
+    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, optimizer_finetune)
+    model = pruner.compress()
+    model.compile(optimizer=optimizer_finetune, loss='sparse_categorical_crossentropy', metrics=['accuracy'])
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--pruner_name', type=str, default='level')
+    parser.add_argument('--batch_size', type=int, default=256)
+    parser.add_argument('--pretrain_epochs', type=int, default=10)
+    parser.add_argument('--prune_epochs', type=int, default=10)
+
+    args = parser.parse_args()
+    main(args)
diff --git a/src/sdk/pynni/nni/compression/tensorflow/__init__.py b/src/sdk/pynni/nni/compression/tensorflow/__init__.py
index 45b6c4e7b8..00d41ee55b 100644
--- a/src/sdk/pynni/nni/compression/tensorflow/__init__.py
+++ b/src/sdk/pynni/nni/compression/tensorflow/__init__.py
@@ -1,6 +1,5 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-from .compressor import LayerInfo, Compressor, Pruner, Quantizer
-from .builtin_pruners import *
-from .builtin_quantizers import *
+from .compressor import Compressor, Pruner
+from .pruning import *
diff --git a/src/sdk/pynni/nni/compression/tensorflow/compressor.py b/src/sdk/pynni/nni/compression/tensorflow/compressor.py
index 542f310f82..e8d73487ac 100644
--- a/src/sdk/pynni/nni/compression/tensorflow/compressor.py
+++ b/src/sdk/pynni/nni/compression/tensorflow/compressor.py
@@ -4,43 +4,28 @@
 import logging
 import tensorflow as tf
 from . import default_layers
-tf.config.experimental_run_functions_eagerly(True)
 
 _logger = logging.getLogger(__name__)
 
 
 class LayerInfo:
-    def __init__(self, keras_layer):
-        self.keras_layer = keras_layer
-        self.name = keras_layer.name
-        self.type = default_layers.get_op_type(type(keras_layer))
-        self.weight_index = default_layers.get_weight_index(self.type)
-        if self.weight_index is not None:
-            self.weight = keras_layer.weights[self.weight_index]
-        self._call = None
-
-def _setattr(model, name, module):
-    name_list = name.split('.')
-    for name in name_list[:-1]:
-        model = getattr(model, name)
-    setattr(model, name_list[-1], module)
+    def __init__(self, layer):
+        self.module = layer
+        self.name = layer._name
+        self.type = type(layer).__name__
+
+
+def _wrap_model(model, wrapped_layers):
+    for key, value in model.__dict__.items():
+        if isinstance(value, tf.keras.Model):
+            _wrap_model(value, wrapped_layers)
+        for layer in wrapped_layers:
+            if value is layer.module:
+                setattr(model, key, layer)
 
-class Compressor:
-    """
-    Abstract base TensorFlow compressor
-    """
 
+class Compressor:
     def __init__(self, model, config_list, optimizer=None):
-        """
-        Record necessary info in class members
-
-        Parameters
-        ----------
-        model : keras model
-            the model user wants to compress
-        config_list : list
-            the configurations that users specify for compression
-        """
         assert isinstance(model, tf.keras.Model)
         self.validate_config(model, config_list)
 
@@ -50,10 +35,6 @@ def __init__(self, model, config_list, optimizer=None):
 
         self.modules_to_compress = None
         self.modules_wrapper = []
-        self.is_wrapped = False
-
-        self._fwd_hook_handles = {}
-        self._fwd_hook_id = 0
 
         for layer, config in self._detect_modules_to_compress():
             wrapper = self._wrap_modules(layer, config)
@@ -61,14 +42,9 @@ def __init__(self, model, config_list, optimizer=None):
         if not self.modules_wrapper:
             _logger.warning('Nothing is configured to compress, please check your model and config list')
 
-        self._wrap_model()
+        _wrap_model(model, self.modules_wrapper)
 
     def _detect_modules_to_compress(self):
-        """
-        detect all modules should be compressed, and save the result in `self.modules_to_compress`.
-
-        The model will be instrumented and user should never edit it after calling this method.
-        """
         if self.modules_to_compress is None:
             self.modules_to_compress = []
             for keras_layer in self.bound_model.layers:
@@ -78,128 +54,67 @@ def _detect_modules_to_compress(self):
                     self.modules_to_compress.append((layer, config))
         return self.modules_to_compress
 
-    def _wrap_model(self):
-        for wrapper in reversed(self.modules_wrapper):
-            _setattr(self.bound_model, wrapper.name, wrapper)
-        self.is_wrapped = True
-
     def compress(self):
-        """
-        Compress the model with algorithm implemented by subclass.
-
-        The model will be instrumented and user should never edit it after calling this method.
-        `self.modules_to_compress` records all the to-be-compressed layers
-        """
-        #modules_to_compress = self.detect_modules_to_compress()
-        #for layer, config in modules_to_compress:
-        #    self._instrument_layer(layer, config)
         return self.bound_model
 
+    def set_wrappers_attribute(self, name, value):
+        for wrapper in self.get_modules_wrapper():
+            setattr(wrapper, name, value)
+
     def get_modules_to_compress(self):
-        """
-        To obtain all the to-be-compressed layers.
-
-        Returns
-        -------
-        self.modules_to_compress : list
-            a list of the layers, each of which is a tuple (`layer`, `config`),
-            `layer` is `LayerInfo`, `config` is a `dict`
-        """
         return self.modules_to_compress
 
     def select_config(self, layer):
-        """
-        Find the configuration for `layer` by parsing `self.config_list`
-
-        Parameters
-        ----------
-        layer: LayerInfo
-            one layer
-
-        Returns
-        -------
-        ret : config or None
-            the retrieved configuration for this layer, if None, this layer should
-            not be compressed
-        """
         ret = None
         if layer.type is None:
             return None
         for config in self.config_list:
             config = config.copy()
-            config['op_types'] = self._expand_config_op_types(config)
-            if layer.type not in config['op_types']:
+            if 'op_types' in config and 'default' in config['op_type']:
+                expanded_op_types = []
+            for op_type in config['op_types']:
+                if op_type == 'default':
+                    expanded_op_types.extend(default_layers.weighted_modules)
+                else:
+                    expanded_op_types.append(op_type)
+            config['op_types'] = expanded_op_types
+
+            if 'op_types' in config and layer.type not in config['op_types']:
                 continue
-            if config.get('op_names') and layer.name not in config['op_names']:
+            if 'op_names' in config and layer.name not in config['op_names']:
                 continue
+
             ret = config
-        if ret is None or ret.get('exclude'):
+
+        if ret is None or 'exclude' is ret:
             return None
         return ret
 
-    def update_epoch(self, epoch):
-        """
-        If user want to update model every epoch, user can override this method.
-        This method should be called at the beginning of each epoch
-
-        Parameters
-        ----------
-        epoch : num
-            the current epoch number
-        """
 
-    def step(self):
-        """
-        If user want to update mask every step, user can override this method
-        """
+    def update_epoch(self, epoch):
+        pass
 
 
     def _wrap_modules(self, layer, config):
         raise NotImplementedError()
 
-    def _instrument_layer(self, layer, config):
-        """
-        This method is implemented in the subclasses, i.e., `Pruner` and `Quantizer`
-
-        Parameters
-        ----------
-        layer : LayerInfo
-            the layer to instrument the compression operation
-        config : dict
-            the configuration for compressing this layer
-        """
-        raise NotImplementedError()
-
-    def _expand_config_op_types(self, config):
-        if config is None:
-            return []
-        op_types = []
 
-        for op_type in config.get('op_types', []):
-            if op_type == 'default':
-                op_types.extend(default_layers.default_layers)
-            else:
-                op_types.append(op_type)
-        return op_types
+    def patch_optimizer(self, **tasks):
+        pass
 
 
 class Pruner(Compressor):
-    """
-    Abstract base TensorFlow pruner
-    """
-
     def __init__(self, model, config_list, optimizer=None):
         super().__init__(model, config_list, optimizer)
         if optimizer is not None:
-            raise RuntimeError('Optimizer patching not implemented yet')
-            #self.patch_optimizer(self.update_mask)
+            self.patch_optimizer(self.update_mask)
 
     def compress(self):
         self.update_mask()
         return self.bound_model
 
     def update_mask(self):
-        for wrapper_idx, wrapper in enumerate(self.modules_wrapper):
+        for wrapper_idx, wrapper in enumerate(self.get_modules_wrapper()):
             masks = self.calc_mask(wrapper, wrapper_idx=wrapper_idx)
             if masks is not None:
                 for k in masks:
@@ -207,38 +122,29 @@ def update_mask(self):
                     setattr(wrapper, k, masks[k])
 
     def calc_mask(self, wrapper, **kwargs):
-        """
-        Pruners should overload this method to provide mask for weight tensors.
-        The mask must have the same shape and type comparing to the weight.
-        It will be applied with `mul()` operation on the weight.
-        This method is effectively hooked to `forward()` method of the model.
-
-        Parameters
-        ----------
-        layer : LayerInfo
-            calculate mask for `layer`'s weight
-        config : dict
-            the configuration for generating the mask
-        """
         raise NotImplementedError("Pruners must overload calc_mask()")
 
     def _wrap_modules(self, layer, config):
-        """
-        Create a wrapper forward function to replace the original one.
-
-        Parameters
-        ----------
-        layer : LayerInfo
-            the layer to instrument the mask
-        config : dict
-            the configuration for generating the mask
-        """
+        _logger.info('Module detected to compress : %s.', layer.name)
         return PrunerModuleWrapper(layer.module, layer.name, layer.type, config, self)
 
-#class Quantizer(Compressor):
-#    """
-#    Abstract base TensorFlow quantizer
-#    """
-#
-#    def quantize_weight(self, weight, config, op, op_type, op_name):
-#        raise NotImplementedError("Quantizer must overload quantize_weight()")
+
+class PrunerModuleWrapper(tf.keras.Model):
+    def __init__(self, module, module_name, module_type, config, pruner):
+        super().__init__()
+        self.module = module
+        self.name = module_name
+        self.type = module_type
+        self.config = config
+        self.pruner = pruner
+        self.masks = []
+        for weight in module.weights:
+            self.masks.append(tf.ones_like(weight))
+            # TODO: filter weight name like 'kernel'/'bias'/etc?
+
+    def call(self, *inputs):
+        new_weights = []
+        for mask, weight in zip(self.masks, self.module.weights):
+            new_weights.append(tf.math.multiply(mask, weight).numpy())
+        self.module.set_weights(new_weights)
+        return self.module(*inputs)
diff --git a/src/sdk/pynni/nni/compression/tensorflow/default_layers.py b/src/sdk/pynni/nni/compression/tensorflow/default_layers.py
index 2ecc46e3e3..4d7e6d8aed 100644
--- a/src/sdk/pynni/nni/compression/tensorflow/default_layers.py
+++ b/src/sdk/pynni/nni/compression/tensorflow/default_layers.py
@@ -1,31 +1,9 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-from tensorflow import keras
-
-supported_layers = {
-    keras.layers.Conv1D: ('Conv1D', 0),
-    keras.layers.Conv2D: ('Conv2D', 0),
-    keras.layers.Conv2DTranspose: ('Conv2DTranspose', 0),
-    keras.layers.Conv3D: ('Conv3D', 0),
-    keras.layers.Conv3DTranspose: ('Conv3DTranspose', 0),
-    keras.layers.ConvLSTM2D: ('ConvLSTM2D', 0),
-    keras.layers.Dense: ('Dense', 0),
-    keras.layers.Embedding: ('Embedding', 0),
-    keras.layers.GRU: ('GRU', 0),
-    keras.layers.LSTM: ('LSTM', 0),
-}
-
-default_layers = [x[0] for x in supported_layers.values()]
-
-def get_op_type(layer_type):
-    if layer_type in supported_layers:
-        return supported_layers[layer_type][0]
-    else:
-        return None
-
-def get_weight_index(op_type):
-    for k in supported_layers:
-        if supported_layers[k][0] == op_type:
-            return supported_layers[k][1]
-    return None
+weighted_modules = [
+    'Conv1d', 'Conv2d', 'Conv3d', 'ConvTranspose1d', 'ConvTranspose2d', 'ConvTranspose3d',
+    'Linear', 'Bilinear',
+    'PReLU',
+    'Embedding', 'EmbeddingBag',
+]
diff --git a/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py b/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py
index 0496b34983..c4841798ba 100644
--- a/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py
+++ b/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py
@@ -12,28 +12,21 @@ def __init__(self, model, config_list, pruning_algorithm='level', optimizer=None
         self.set_wrapper_attribute('if_calculated', False)
         self.masker = MASKER_DICT[pruning_alogrithm](model, self, **algo_kwargs)
 
+
     def validate_config(self, model, config_list):
-        pass
+        pass  # TODO
+
 
     def calc_mask(self, wrapper, wrapper_idx=None):
         if wrapper.if_calculated:
             return None
-
         sparsity = wrapper.config['sparsity']
-        if not wrapper.if_calculated:
-            masks = self.masker.calc_mask(sparsity=sparsity, wrapper=wrapper, wrapper_idx=wrapper_idx)
-
-            if masks is not None:
-                wrapper.if_calculated = True
-            return masks
-        else:
-            return None
+        masks = self.masker.calc_mask(sparsity=sparsity, wrapper=wrapper, wrapper_idx=wrapper_idx)
+        if masks is not None:
+            wrapper.if_calculated = True
+        return masks
 
 
-MASKER_DICT = {
-    'level': LevelPrunerMasker,
-}
-
 class WeightMasker:
     def __init__(self, model, pruner, **kwargs):
         self.model = model
@@ -42,6 +35,7 @@ def __init__(self, model, pruner, **kwargs):
     def calc_mask(self, sparsity, wrapper, wrapper_idx=None):
         raise NotImplementedError()
 
+
 class LevelPrunerMasker(WeightMasker):
     def calc_mask(self, sparsity, wrapper, wrapper_idx=None):
         weight = wrapper.module.weight * wrapper.weight_mask
@@ -51,3 +45,8 @@ def calc_mask(self, sparsity, wrapper, wrapper_idx=None):
         threshold = tf.reduce_max(tf.topk(tf.reshape(w_abs, [-1]), k, largest=False)[0])
         mask_weight = tf.cast((w_abs > threshold), weight.dtype)
         return {'weight_mask': mask_weight}
+
+
+MASKER_DICT = {
+    'level': LevelPrunerMasker,
+}

From 0670561fb6c0f9f1f4b12a6c58f80c5379025ff0 Mon Sep 17 00:00:00 2001
From: liuzhe <zhliu1@microsoft.com>
Date: Wed, 5 Aug 2020 10:42:04 +0800
Subject: [PATCH 3/9] fix default layer names

---
 src/sdk/pynni/nni/compression/tensorflow/default_layers.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/sdk/pynni/nni/compression/tensorflow/default_layers.py b/src/sdk/pynni/nni/compression/tensorflow/default_layers.py
index 4d7e6d8aed..0c729bd883 100644
--- a/src/sdk/pynni/nni/compression/tensorflow/default_layers.py
+++ b/src/sdk/pynni/nni/compression/tensorflow/default_layers.py
@@ -2,8 +2,8 @@
 # Licensed under the MIT license.
 
 weighted_modules = [
-    'Conv1d', 'Conv2d', 'Conv3d', 'ConvTranspose1d', 'ConvTranspose2d', 'ConvTranspose3d',
-    'Linear', 'Bilinear',
+    'Conv1D', 'Conv2D', 'Conv3D', 'Conv1DTranspose', 'Conv2DTranspose', 'Conv3DTranspose',
+    'Dense',
     'PReLU',
-    'Embedding', 'EmbeddingBag',
+    'Embedding',
 ]

From 515121f992bb48fecd40b2b6de562c293e7f3355 Mon Sep 17 00:00:00 2001
From: liuzhe <zhliu1@microsoft.com>
Date: Fri, 7 Aug 2020 16:36:27 +0800
Subject: [PATCH 4/9] refactor

---
 examples/model_compress/model_prune_tf.py     |   7 +-
 .../nni/compression/tensorflow/compressor.py  | 223 +++++++++---------
 .../tensorflow/pruning/one_shot.py            |  55 +++--
 3 files changed, 155 insertions(+), 130 deletions(-)

diff --git a/examples/model_compress/model_prune_tf.py b/examples/model_compress/model_prune_tf.py
index 3a928bc6da..9e672aff49 100644
--- a/examples/model_compress/model_prune_tf.py
+++ b/examples/model_compress/model_prune_tf.py
@@ -44,10 +44,10 @@ def create_model(model_name='naive'):
     ])
 
 
-def create_pruner(model, pruner_name, optimizer=None):
+def create_pruner(model, pruner_name):
     pruner_class = prune_config[pruner_name]['pruner_class']
     config_list = prune_config[pruner_name]['config_list']
-    return pruner_class(model, config_list, optimizer)
+    return pruner_class(model, config_list)
 
 
 def main(args):
@@ -64,9 +64,10 @@ def main(args):
 
     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, optimizer_finetune)
+    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)
 
 
 if __name__ == '__main__':
diff --git a/src/sdk/pynni/nni/compression/tensorflow/compressor.py b/src/sdk/pynni/nni/compression/tensorflow/compressor.py
index e8d73487ac..72d65077f8 100644
--- a/src/sdk/pynni/nni/compression/tensorflow/compressor.py
+++ b/src/sdk/pynni/nni/compression/tensorflow/compressor.py
@@ -9,142 +9,151 @@
 
 
 class LayerInfo:
-    def __init__(self, layer):
-        self.module = layer
-        self.name = layer._name
+    def __init__(self, layer, path=None):
+        self.layer = layer
+        self.name = layer.name
         self.type = type(layer).__name__
-
-
-def _wrap_model(model, wrapped_layers):
-    for key, value in model.__dict__.items():
-        if isinstance(value, tf.keras.Model):
-            _wrap_model(value, wrapped_layers)
-        for layer in wrapped_layers:
-            if value is layer.module:
-                setattr(model, key, layer)
+        self.path = path
+        self.config = None
 
 
 class Compressor:
-    def __init__(self, model, config_list, optimizer=None):
+    def __init__(self, LayerWrapperClass, model, config_list):
         assert isinstance(model, tf.keras.Model)
         self.validate_config(model, config_list)
 
         self.bound_model = model
-        self.config_list = config_list
-        self.optimizer = optimizer
-
-        self.modules_to_compress = None
-        self.modules_wrapper = []
+        self.wrappers = []
 
-        for layer, config in self._detect_modules_to_compress():
-            wrapper = self._wrap_modules(layer, config)
-            self.modules_wrapper.append(wrapper)
-        if not self.modules_wrapper:
+        for layer_info in _detect_layers_to_compress(model, config_list):
+            self.wrappers.append(LayerWrapperClass(layer_info, self))
+        if not self.wrappers:
             _logger.warning('Nothing is configured to compress, please check your model and config list')
 
-        _wrap_model(model, self.modules_wrapper)
-
-    def _detect_modules_to_compress(self):
-        if self.modules_to_compress is None:
-            self.modules_to_compress = []
-            for keras_layer in self.bound_model.layers:
-                layer = LayerInfo(keras_layer)
-                config = self.select_config(layer)
-                if config is not None:
-                    self.modules_to_compress.append((layer, config))
-        return self.modules_to_compress
-
-    def compress(self):
-        return self.bound_model
+        _instrument_model(model, self.wrappers)
 
     def set_wrappers_attribute(self, name, value):
-        for wrapper in self.get_modules_wrapper():
+        for wrapper in self.wrappers:
             setattr(wrapper, name, value)
 
-    def get_modules_to_compress(self):
-        return self.modules_to_compress
-
-    def select_config(self, layer):
-        ret = None
-        if layer.type is None:
-            return None
-        for config in self.config_list:
-            config = config.copy()
-            if 'op_types' in config and 'default' in config['op_type']:
-                expanded_op_types = []
-            for op_type in config['op_types']:
-                if op_type == 'default':
-                    expanded_op_types.extend(default_layers.weighted_modules)
-                else:
-                    expanded_op_types.append(op_type)
-            config['op_types'] = expanded_op_types
-
-            if 'op_types' in config and layer.type not in config['op_types']:
-                continue
-            if 'op_names' in config and layer.name not in config['op_names']:
-                continue
-
-            ret = config
-
-        if ret is None or 'exclude' is ret:
-            return None
-        return ret
-
-
-    def update_epoch(self, epoch):
-        pass
-
-
-    def _wrap_modules(self, layer, config):
-        raise NotImplementedError()
-
-
-    def patch_optimizer(self, **tasks):
-        pass
-
 
 class Pruner(Compressor):
-    def __init__(self, model, config_list, optimizer=None):
-        super().__init__(model, config_list, optimizer)
-        if optimizer is not None:
-            self.patch_optimizer(self.update_mask)
+    def __init__(self, model, config_list):
+        super().__init__(PrunerLayerWrapper, model, config_list)
+        #self.callback = PrunerCallback(self)
 
     def compress(self):
         self.update_mask()
         return self.bound_model
 
     def update_mask(self):
-        for wrapper_idx, wrapper in enumerate(self.get_modules_wrapper()):
-            masks = self.calc_mask(wrapper, wrapper_idx=wrapper_idx)
+        for wrapper_idx, wrapper in enumerate(self.wrappers):
+            masks = self.calc_masks(wrapper, wrapper_idx=wrapper_idx)
             if masks is not None:
-                for k in masks:
-                    assert hasattr(wrapper, k)
-                    setattr(wrapper, k, masks[k])
-
-    def calc_mask(self, wrapper, **kwargs):
-        raise NotImplementedError("Pruners must overload calc_mask()")
+                wrapper.masks = masks
 
-    def _wrap_modules(self, layer, config):
-        _logger.info('Module detected to compress : %s.', layer.name)
-        return PrunerModuleWrapper(layer.module, layer.name, layer.type, config, self)
+    def calc_masks(self, wrapper, **kwargs):
+        # TODO: maybe it should be able to calc on weight-granularity, beside from layer-granularity
+        raise NotImplementedError("Pruners must overload calc_masks()")
 
 
-class PrunerModuleWrapper(tf.keras.Model):
-    def __init__(self, module, module_name, module_type, config, pruner):
+class PrunerLayerWrapper(tf.keras.Model):
+    def __init__(self, layer_info, pruner):
         super().__init__()
-        self.module = module
-        self.name = module_name
-        self.type = module_type
-        self.config = config
+        self.layer_info = layer_info
+        self.layer = layer_info.layer
+        self.config = layer_info.config
         self.pruner = pruner
-        self.masks = []
-        for weight in module.weights:
-            self.masks.append(tf.ones_like(weight))
-            # TODO: filter weight name like 'kernel'/'bias'/etc?
+        self.masks = {}
+        _logger.info('Layer detected to compress: %s', self.layer.name)
 
     def call(self, *inputs):
         new_weights = []
-        for mask, weight in zip(self.masks, self.module.weights):
-            new_weights.append(tf.math.multiply(mask, weight).numpy())
-        self.module.set_weights(new_weights)
-        return self.module(*inputs)
+        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())
+            else:
+                new_weights.append(weight.numpy())
+        self.layer.set_weights(new_weights)
+        return self.layer(*inputs)
+
+
+# TODO: designed to replace `patch_optimizer`
+#class PrunerCallback(tf.keras.callbacks.Callback):
+#    def __init__(self, pruner):
+#        super().__init__()
+#        self._pruner = pruner
+#
+#    def on_train_batch_end(self, batch, logs=None):
+#        self._pruner.update_mask()
+
+
+def _detect_layers_to_compress(model, config_list):
+    located_layers = _locate_layers(model)
+    ret = []
+    for layer in model.layers:
+        config = _select_config(LayerInfo(layer), config_list)
+        if config is not None:
+            if id(layer) not in located_layers:
+                _logger.error('Failed to locate layer %s in model. The layer will not be compressed. '
+                              'This is a bug in NNI, feel free to fire an issue.', layer.name)
+                continue
+            layer_info = located_layers[id(layer)]
+            layer_info.config = config
+            ret.append(layer_info)
+    return ret
+
+def _locate_layers(model, cur_path=[]):
+    # FIXME: this cannot find layers contained in list, dict, non-model custom classes, etc
+    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):
+            if isinstance(item, tf.keras.Model):
+                ret.update(_locate_layers(item, cur_path + [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)))
+    return ret
+
+def _select_config(layer_info, config_list):
+    ret = None
+    for config in config_list:
+        if 'op_types' in config:
+            match = layer_info.type in config['op_types']
+            match_default = 'default' in config['op_types'] and layer_info.type in default_layers.weighted_modules
+            if not match and not match_default:
+                continue
+        if 'op_names' in config and layer_info.name not in config['op_names']:
+            continue
+        ret = config
+    if ret is None or 'exclude' in ret:
+        return None
+    return ret
+
+
+def _instrument_model(model, wrappers):
+    for wrapper in wrappers:
+        cur = model
+        for key in wrapper.layer_info.path[:-1]:
+            if isinstance(key, int):
+                cur = cur[key]
+            else:
+                cur = getattr(cur, key)
+        key = wrapper.layer_info.path[-1]
+        if isinstance(key, int):
+            cur[key] = wrapper
+        else:
+            setattr(cur, key, wrapper)
diff --git a/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py b/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py
index c4841798ba..fd4e768858 100644
--- a/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py
+++ b/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py
@@ -3,48 +3,63 @@
 from ..compressor import Pruner
 
 __all__ = [
-    'OneshotPruner'
+    'OneshotPruner',
+    'LevelPruner',
 ]
 
 class OneshotPruner(Pruner):
-    def __init__(self, model, config_list, pruning_algorithm='level', optimizer=None, **algo_kwargs):
-        super().__init__(model, config_list, optimizer)
-        self.set_wrapper_attribute('if_calculated', False)
-        self.masker = MASKER_DICT[pruning_alogrithm](model, self, **algo_kwargs)
-
+    def __init__(self, model, config_list, pruning_algorithm='level', **algo_kwargs):
+        super().__init__(model, config_list)
+        self.set_wrappers_attribute('calculated', False)
+        self.masker = MASKER_DICT[pruning_algorithm](model, self, **algo_kwargs)
 
     def validate_config(self, model, config_list):
         pass  # TODO
 
-
-    def calc_mask(self, wrapper, wrapper_idx=None):
-        if wrapper.if_calculated:
+    def calc_masks(self, wrapper, wrapper_idx=None):
+        if wrapper.calculated:
             return None
         sparsity = wrapper.config['sparsity']
-        masks = self.masker.calc_mask(sparsity=sparsity, wrapper=wrapper, wrapper_idx=wrapper_idx)
+        masks = self.masker.calc_masks(sparsity, wrapper, wrapper_idx)
         if masks is not None:
-            wrapper.if_calculated = True
+            wrapper.calculated = True
         return masks
 
 
+class LevelPruner(OneshotPruner):
+    def __init__(self, model, config_list):
+        super().__init__(model, config_list, pruning_algorithm='level')
+
+
 class WeightMasker:
     def __init__(self, model, pruner, **kwargs):
         self.model = model
         self.pruner = pruner
 
-    def calc_mask(self, sparsity, wrapper, wrapper_idx=None):
+    def calc_masks(self, sparsity, wrapper, wrapper_idx=None):
         raise NotImplementedError()
 
 
 class LevelPrunerMasker(WeightMasker):
-    def calc_mask(self, sparsity, wrapper, wrapper_idx=None):
-        weight = wrapper.module.weight * wrapper.weight_mask
-        w_abs = tf.abs(w_abs)
-        k = int(tf.size(weight) * sparsity)
-        assert k > 0  # FIXME
-        threshold = tf.reduce_max(tf.topk(tf.reshape(w_abs, [-1]), k, largest=False)[0])
-        mask_weight = tf.cast((w_abs > threshold), weight.dtype)
-        return {'weight_mask': mask_weight}
+    def calc_masks(self, sparsity, wrapper, wrapper_idx=None):
+        masks = {}
+        for i, weight_variable in enumerate(wrapper.layer.weights):
+            if weight_variable.name == 'bias':
+                continue
+
+            k = int(tf.size(weight_variable) * sparsity)
+            if k == 0:
+                continue
+
+            weight = weight_variable.read_value()
+            if wrapper.masks[weight.name] is not None:
+                weight = tf.math.multiply(weight, wrapper.masks[weight.name])
+
+            w_abs = tf.math.abs(tf.reshape(weight, [-1]))
+            threshold = tf.math.topk(w_abs, k)[0][0]
+            mask = tf.math.greater(w_abs, threshold)
+            masks[weight.name] = tf.cast(mask, weight.dtype)
+        return masks
 
 
 MASKER_DICT = {

From 0177ba4296abfc2c130ab9f556990b9ccbc1086b Mon Sep 17 00:00:00 2001
From: liuzhe <zhliu1@microsoft.com>
Date: Fri, 7 Aug 2020 17:21:57 +0800
Subject: [PATCH 5/9] remove old ut

---
 examples/model_compress/model_prune_tf.py     |  1 +
 .../tensorflow/pruning/one_shot.py            | 10 ++--
 src/sdk/pynni/tests/test_compressor.py        | 55 +------------------
 3 files changed, 7 insertions(+), 59 deletions(-)

diff --git a/examples/model_compress/model_prune_tf.py b/examples/model_compress/model_prune_tf.py
index 9e672aff49..9ae5b0dd10 100644
--- a/examples/model_compress/model_prune_tf.py
+++ b/examples/model_compress/model_prune_tf.py
@@ -41,6 +41,7 @@ def create_model(model_name='naive'):
         tf.keras.layers.Dense(units=500),
         tf.keras.layers.ReLU(),
         tf.keras.layers.Dense(units=10),
+        tf.keras.layers.Softmax()
     ])
 
 
diff --git a/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py b/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py
index fd4e768858..feefe95f7d 100644
--- a/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py
+++ b/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py
@@ -47,18 +47,18 @@ def calc_masks(self, sparsity, wrapper, wrapper_idx=None):
             if weight_variable.name == 'bias':
                 continue
 
-            k = int(tf.size(weight_variable) * sparsity)
+            k = int(tf.size(weight_variable).numpy() * sparsity)
             if k == 0:
                 continue
 
             weight = weight_variable.read_value()
-            if wrapper.masks[weight.name] is not None:
-                weight = tf.math.multiply(weight, wrapper.masks[weight.name])
+            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.topk(w_abs, k)[0][0]
+            threshold = tf.math.top_k(w_abs, k)[0][0]
             mask = tf.math.greater(w_abs, threshold)
-            masks[weight.name] = tf.cast(mask, weight.dtype)
+            masks[weight_variable.name] = tf.cast(mask, weight.dtype)
         return masks
 
 
diff --git a/src/sdk/pynni/tests/test_compressor.py b/src/sdk/pynni/tests/test_compressor.py
index 7641ae7d25..4322f22750 100644
--- a/src/sdk/pynni/tests/test_compressor.py
+++ b/src/sdk/pynni/tests/test_compressor.py
@@ -3,33 +3,12 @@
 
 from unittest import TestCase, main
 import numpy as np
-import tensorflow as tf
 import torch
 import torch.nn.functional as F
 import schema
 import nni.compression.torch as torch_compressor
 import math
 
-if tf.__version__ >= '2.0':
-    import nni.compression.tensorflow as tf_compressor
-
-
-def get_tf_model():
-    model = tf.keras.models.Sequential([
-        tf.keras.layers.Conv2D(filters=5, kernel_size=7, input_shape=[28, 28, 1], activation='relu', padding="SAME"),
-        tf.keras.layers.MaxPooling2D(pool_size=2),
-        tf.keras.layers.Conv2D(filters=10, kernel_size=3, activation='relu', padding="SAME"),
-        tf.keras.layers.MaxPooling2D(pool_size=2),
-        tf.keras.layers.Flatten(),
-        tf.keras.layers.Dense(units=128, activation='relu'),
-        tf.keras.layers.Dropout(0.5),
-        tf.keras.layers.Dense(units=10, activation='softmax'),
-    ])
-    model.compile(loss="sparse_categorical_crossentropy",
-                  optimizer=tf.keras.optimizers.SGD(lr=1e-3),
-                  metrics=["accuracy"])
-    return model
-
 
 class TorchModel(torch.nn.Module):
     def __init__(self):
@@ -52,13 +31,6 @@ def forward(self, x):
         return F.log_softmax(x, dim=1)
 
 
-def tf2(func):
-    def test_tf2_func(*args):
-        if tf.__version__ >= '2.0':
-            func(*args)
-
-    return test_tf2_func
-
 class CompressorTestCase(TestCase):
     def test_torch_quantizer_modules_detection(self):
         # test if modules can be detected
@@ -92,11 +64,6 @@ def test_torch_level_pruner(self):
         configure_list = [{'sparsity': 0.8, 'op_types': ['default']}]
         torch_compressor.LevelPruner(model, configure_list, optimizer).compress()
 
-    @tf2
-    def test_tf_level_pruner(self):
-        configure_list = [{'sparsity': 0.8, 'op_types': ['default']}]
-        tf_compressor.LevelPruner(get_tf_model(), configure_list).compress()
-
     def test_torch_naive_quantizer(self):
         model = TorchModel()
         configure_list = [{
@@ -108,10 +75,6 @@ def test_torch_naive_quantizer(self):
         }]
         torch_compressor.NaiveQuantizer(model, configure_list).compress()
 
-    @tf2
-    def test_tf_naive_quantizer(self):
-        tf_compressor.NaiveQuantizer(get_tf_model(), [{'op_types': ['default']}]).compress()
-
     def test_torch_fpgm_pruner(self):
         """
         With filters(kernels) weights defined as above (w), it is obvious that w[4] and w[5] is the Geometric Median
@@ -141,23 +104,7 @@ def test_torch_fpgm_pruner(self):
         masks = pruner.calc_mask(model.conv2)
         assert all(torch.sum(masks['weight_mask'], (1, 2, 3)).numpy() == np.array([125., 125., 0., 0., 0., 0., 0., 0., 125., 125.]))
 
-    @tf2
-    def test_tf_fpgm_pruner(self):
-        w = np.array([np.ones((5, 3, 3)) * (i+1) for i in range(10)]).astype(np.float32)
-        model = get_tf_model()
-        config_list = [{'sparsity': 0.2, 'op_types': ['Conv2D']}]
-
-        pruner = tf_compressor.FPGMPruner(model, config_list)
-        weights = model.layers[2].weights
-        weights[0] = np.array(w).astype(np.float32).transpose([2, 3, 0, 1]).transpose([0, 1, 3, 2])
-        model.layers[2].set_weights([weights[0], weights[1].numpy()])
-
-        layer = tf_compressor.compressor.LayerInfo(model.layers[2])
-        masks = pruner.calc_mask(layer, config_list[0]).numpy()
-        masks = masks.reshape((-1, masks.shape[-1])).transpose([1, 0])
-
-        assert all(masks.sum((1)) == np.array([45., 45., 45., 45., 0., 0., 45., 45., 45., 45.]))
-        
+       
     def test_torch_l1filter_pruner(self):
         """
         Filters with the minimum sum of the weights' L1 norm are pruned in this paper:

From 949f52b50b04011d5bee986a21c6357d5224f94e Mon Sep 17 00:00:00 2001
From: liuzhe <zhliu1@microsoft.com>
Date: Fri, 7 Aug 2020 19:26:01 +0800
Subject: [PATCH 6/9] add base classes' doc

---
 examples/model_compress/model_prune_tf.py     |   2 +-
 .../compression/tensorflow/builtin_pruners.py | 195 ------------------
 .../tensorflow/builtin_quantizers.py          |  74 -------
 .../nni/compression/tensorflow/compressor.py  | 157 +++++++++++++-
 ...compressor.py => test_compressor_torch.py} |   0
 5 files changed, 150 insertions(+), 278 deletions(-)
 delete mode 100644 src/sdk/pynni/nni/compression/tensorflow/builtin_pruners.py
 delete mode 100644 src/sdk/pynni/nni/compression/tensorflow/builtin_quantizers.py
 rename src/sdk/pynni/tests/{test_compressor.py => test_compressor_torch.py} (100%)

diff --git a/examples/model_compress/model_prune_tf.py b/examples/model_compress/model_prune_tf.py
index 9ae5b0dd10..99e8278df4 100644
--- a/examples/model_compress/model_prune_tf.py
+++ b/examples/model_compress/model_prune_tf.py
@@ -10,7 +10,7 @@
         'model_name': 'naive',
         'pruner_class': nni.compression.tensorflow.LevelPruner,
         'config_list': [{
-            'sparsity': 0.5,
+            'sparsity': 0.9,
             'op_types': ['default'],
         }]
     },
diff --git a/src/sdk/pynni/nni/compression/tensorflow/builtin_pruners.py b/src/sdk/pynni/nni/compression/tensorflow/builtin_pruners.py
deleted file mode 100644
index 9ff3d71b92..0000000000
--- a/src/sdk/pynni/nni/compression/tensorflow/builtin_pruners.py
+++ /dev/null
@@ -1,195 +0,0 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT license.
-
-import logging
-import numpy as np
-import tensorflow as tf
-from .compressor import Pruner
-
-__all__ = ['LevelPruner', 'AGP_Pruner', 'FPGMPruner']
-
-_logger = logging.getLogger(__name__)
-
-
-class LevelPruner(Pruner):
-    def __init__(self, model, config_list):
-        """
-        config_list: supported keys:
-            - sparsity
-        """
-        super().__init__(model, config_list)
-        self.mask_list = {}
-        self.if_init_list = {}
-
-    def calc_mask(self, layer, config):
-        weight = layer.weight
-        op_name = layer.name
-        if self.if_init_list.get(op_name, True):
-            threshold = tf.contrib.distributions.percentile(tf.abs(weight), config['sparsity'] * 100)
-            mask = tf.cast(tf.math.greater(tf.abs(weight), threshold), weight.dtype)
-            self.mask_list.update({op_name: mask})
-            self.if_init_list.update({op_name: False})
-        else:
-            mask = self.mask_list[op_name]
-        return mask
-
-
-class AGP_Pruner(Pruner):
-    """An automated gradual pruning algorithm that prunes the smallest magnitude
-    weights to achieve a preset level of network sparsity.
-    Michael Zhu and Suyog Gupta, "To prune, or not to prune: exploring the
-    efficacy of pruning for model compression", 2017 NIPS Workshop on Machine
-    Learning of Phones and other Consumer Devices,
-    https://arxiv.org/pdf/1710.01878.pdf
-    """
-
-    def __init__(self, model, config_list):
-        """
-        config_list: supported keys:
-            - initial_sparsity
-            - final_sparsity: you should make sure initial_sparsity <= final_sparsity
-            - start_epoch: start epoch numer begin update mask
-            - end_epoch: end epoch number stop update mask
-            - frequency: if you want update every 2 epoch, you can set it 2
-        """
-        super().__init__(model, config_list)
-        self.mask_list = {}
-        self.if_init_list = {}
-        self.now_epoch = tf.Variable(0)
-        self.assign_handler = []
-
-    def calc_mask(self, layer, config):
-        weight = layer.weight
-        op_name = layer.name
-        start_epoch = config.get('start_epoch', 0)
-        freq = config.get('frequency', 1)
-        if self.now_epoch >= start_epoch and self.if_init_list.get(op_name, True) and (
-                self.now_epoch - start_epoch) % freq == 0:
-            target_sparsity = self.compute_target_sparsity(config)
-            threshold = tf.contrib.distributions.percentile(weight, target_sparsity * 100)
-            # stop gradient in case gradient change the mask
-            mask = tf.stop_gradient(tf.cast(tf.math.greater(weight, threshold), weight.dtype))
-            self.assign_handler.append(tf.assign(weight, weight * mask))
-            self.mask_list.update({op_name: tf.constant(mask)})
-            self.if_init_list.update({op_name: False})
-        else:
-            mask = self.mask_list[op_name]
-        return mask
-
-    def compute_target_sparsity(self, config):
-        end_epoch = config.get('end_epoch', 1)
-        start_epoch = config.get('start_epoch', 0)
-        freq = config.get('frequency', 1)
-        final_sparsity = config.get('final_sparsity', 0)
-        initial_sparsity = config.get('initial_sparsity', 0)
-
-        if end_epoch <= start_epoch or initial_sparsity >= final_sparsity:
-            _logger.warning('your end epoch <= start epoch or initial_sparsity >= final_sparsity')
-            return final_sparsity
-
-        now_epoch = tf.minimum(self.now_epoch, tf.constant(end_epoch))
-        span = int(((end_epoch - start_epoch - 1) // freq) * freq)
-        assert span > 0
-        base = tf.cast(now_epoch - start_epoch, tf.float32) / span
-        target_sparsity = (final_sparsity +
-                           (initial_sparsity - final_sparsity) *
-                           (tf.pow(1.0 - base, 3)))
-        return target_sparsity
-
-    def update_epoch(self, epoch, sess):
-        sess.run(self.assign_handler)
-        sess.run(tf.assign(self.now_epoch, int(epoch)))
-        for k in self.if_init_list:
-            self.if_init_list[k] = True
-
-class FPGMPruner(Pruner):
-    """
-    A filter pruner via geometric median.
-    "Filter Pruning via Geometric Median for Deep Convolutional Neural Networks Acceleration",
-    https://arxiv.org/pdf/1811.00250.pdf
-    """
-
-    def __init__(self, model, config_list):
-        """
-        Parameters
-        ----------
-        model : pytorch model
-            the model user wants to compress
-        config_list: list
-            support key for each list item:
-                - sparsity: percentage of convolutional filters to be pruned.
-        """
-        super().__init__(model, config_list)
-        self.mask_dict = {}
-        self.assign_handler = []
-        self.epoch_pruned_layers = set()
-
-    def calc_mask(self, layer, config):
-        """
-        Supports Conv1D, Conv2D
-        filter dimensions for Conv1D:
-        LEN: filter length
-        IN: number of input channel
-        OUT: number of output channel
-
-        filter dimensions for Conv2D:
-        H: filter height
-        W: filter width
-        IN: number of input channel
-        OUT: number of output channel
-
-        Parameters
-        ----------
-        layer : LayerInfo
-            calculate mask for `layer`'s weight
-        config : dict
-            the configuration for generating the mask
-        """
-
-        weight = layer.weight
-        op_type = layer.type
-        op_name = layer.name
-        assert 0 <= config.get('sparsity') < 1
-        assert op_type in ['Conv1D', 'Conv2D']
-        assert op_type in config['op_types']
-
-        if layer.name in self.epoch_pruned_layers:
-            assert layer.name in self.mask_dict
-            return self.mask_dict.get(layer.name)
-
-        try:
-            w = tf.stop_gradient(tf.transpose(tf.reshape(weight, (-1, weight.shape[-1])), [1, 0]))
-            masks = np.ones(w.shape)
-            num_filters = w.shape[0]
-            num_prune = int(num_filters * config.get('sparsity'))
-            if num_filters < 2 or num_prune < 1:
-                return masks
-            min_gm_idx = self._get_min_gm_kernel_idx(w, num_prune)
-
-            for idx in min_gm_idx:
-                masks[idx] = 0.
-        finally:
-            masks = tf.reshape(tf.transpose(masks, [1, 0]), weight.shape)
-            masks = tf.Variable(masks)
-            self.mask_dict.update({op_name: masks})
-            self.epoch_pruned_layers.add(layer.name)
-
-        return masks
-
-    def _get_min_gm_kernel_idx(self, weight, n):
-        dist_list = []
-        for out_i in range(weight.shape[0]):
-            dist_sum = self._get_distance_sum(weight, out_i)
-            dist_list.append((dist_sum, out_i))
-        min_gm_kernels = sorted(dist_list, key=lambda x: x[0])[:n]
-        return [x[1] for x in min_gm_kernels]
-
-    def _get_distance_sum(self, weight, out_idx):
-        anchor_w = tf.tile(tf.expand_dims(weight[out_idx], 0), [weight.shape[0], 1])
-        x = weight - anchor_w
-        x = tf.math.reduce_sum((x*x), -1)
-        x = tf.math.sqrt(x)
-        return tf.math.reduce_sum(x)
-
-    def update_epoch(self, epoch):
-        self.epoch_pruned_layers = set()
diff --git a/src/sdk/pynni/nni/compression/tensorflow/builtin_quantizers.py b/src/sdk/pynni/nni/compression/tensorflow/builtin_quantizers.py
deleted file mode 100644
index 3f54cbfb12..0000000000
--- a/src/sdk/pynni/nni/compression/tensorflow/builtin_quantizers.py
+++ /dev/null
@@ -1,74 +0,0 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT license.
-
-import logging
-import tensorflow as tf
-from .compressor import Quantizer
-
-__all__ = ['NaiveQuantizer', 'QAT_Quantizer', 'DoReFaQuantizer']
-
-_logger = logging.getLogger(__name__)
-
-
-class NaiveQuantizer(Quantizer):
-    """quantize weight to 8 bits
-    """
-    def __init__(self, model, config_list):
-        super().__init__(model, config_list)
-        self.layer_scale = {}
-
-    def quantize_weight(self, weight, config, op_name, **kwargs):
-        new_scale = tf.reduce_max(tf.abs(weight)) / 127
-        scale = tf.maximum(self.layer_scale.get(op_name, tf.constant(0.0)), new_scale)
-        self.layer_scale[op_name] = scale
-        orig_type = weight.dtype
-        return tf.cast(tf.cast(weight / scale, tf.int8), orig_type) * scale
-
-
-class QAT_Quantizer(Quantizer):
-    """Quantizer using the Quantization and Training scheme, as defined in:
-    Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference
-    http://openaccess.thecvf.com/content_cvpr_2018/papers/Jacob_Quantization_and_Training_CVPR_2018_paper.pdf
-    """
-    def __init__(self, model, config_list):
-        """
-        config_list: supported keys:
-            - q_bits
-        """
-        super().__init__(model, config_list)
-
-    def quantize_weight(self, weight, config, **kwargs):
-        a = tf.stop_gradient(tf.reduce_min(weight))
-        b = tf.stop_gradient(tf.reduce_max(weight))
-        n = tf.cast(2 ** config['q_bits'], tf.float32)
-        scale = b-a/(n-1)
-
-        # use gradient_override_map to change round to idetity for gradient
-        with tf.get_default_graph().gradient_override_map({'Round': 'Identity'}):
-            qw = tf.round((weight-a)/scale)*scale +a
-
-        return qw
-
-
-class DoReFaQuantizer(Quantizer):
-    """Quantizer using the DoReFa scheme, as defined in:
-    Zhou et al., DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks with Low Bitwidth Gradients
-    (https://arxiv.org/abs/1606.06160)
-    """
-    def __init__(self, model, config_list):
-        """
-        config_list: supported keys:
-            - q_bits
-        """
-        super().__init__(model, config_list)
-
-    def quantize_weight(self, weight, config, **kwargs):
-        a = tf.math.tanh(weight)
-        b = a/(2*tf.reduce_max(tf.abs(weight))) + 0.5
-
-        scale = pow(2, config['q_bits'] - 1)
-        # use gradient_override_map to change round to idetity for gradient
-        with tf.get_default_graph().gradient_override_map({'Round': 'Identity'}):
-            qw = tf.round(b*scale)/scale
-        r_qw = 2 * qw - 1
-        return r_qw
diff --git a/src/sdk/pynni/nni/compression/tensorflow/compressor.py b/src/sdk/pynni/nni/compression/tensorflow/compressor.py
index 72d65077f8..6d9679e82d 100644
--- a/src/sdk/pynni/nni/compression/tensorflow/compressor.py
+++ b/src/sdk/pynni/nni/compression/tensorflow/compressor.py
@@ -1,6 +1,10 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+"""
+Abstract base classes for TensorFlow model compression.
+"""
+
 import logging
 import tensorflow as tf
 from . import default_layers
@@ -9,6 +13,32 @@
 
 
 class LayerInfo:
+    """
+    This structure contains all infomation needed to compress a TensorFlow ``Layer``.
+
+
+    Attributes
+    ----------
+    layer : tf.keras.layers.Layer
+        The layer.
+    name : str
+        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
+        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`.
+    config : JSON object
+        Selected configuration for this layer. The format is detailed in tutorial.
+
+    Parameters
+    ----------
+    layer : tf.keras.layers.Layer
+        See attributes section.
+    path : list of str/int
+        See attributes section.
+    """
+
     def __init__(self, layer, path=None):
         self.layer = layer
         self.name = layer.name
@@ -18,6 +48,31 @@ def __init__(self, layer, path=None):
 
 
 class Compressor:
+    """
+    Common base class for all compressors.
+
+    This class is designed for other base classes.
+    Algorithms should inherit ``Pruner`` or Quantizer instead.
+
+
+    Attributes
+    ----------
+    bound_model : tf.keras.Model
+        Compressed user model.
+    wrappers : list of tf.keras.Model
+        A wrapper is an instrumented TF ``Layer``, in ``Model`` format.
+        The list is ordered by preorder traversal.
+
+    Parameters
+    ----------
+    LayerWrapperClass : a class derive from Model
+        The class used to instrument layers.
+    model : tf.keras.Model
+        The user model to be compressed.
+    config_list : list of JSON object
+        User configuration. The format is detailed in tutorial.
+    """
+
     def __init__(self, LayerWrapperClass, model, config_list):
         assert isinstance(model, tf.keras.Model)
         self.validate_config(model, config_list)
@@ -33,31 +88,108 @@ def __init__(self, LayerWrapperClass, model, config_list):
         _instrument_model(model, self.wrappers)
 
     def set_wrappers_attribute(self, name, value):
+        """
+        Call ``setattr`` on all wrappers.
+        """
         for wrapper in self.wrappers:
             setattr(wrapper, name, value)
 
 
 class Pruner(Compressor):
+    """
+    Base class for pruning algorithms.
+
+    End users should use ``compress`` and callback APIs to prune their models.
+
+    The underlying model is instrumented upon initialization of pruner object.
+    So if you want to pre-train the model, train it before creating pruner object.
+
+    The compressed model can only execute in eager mode.
+
+    Algorithm developers should override ``calc_masks`` method to specify pruning strategy.
+
+    Parameters
+    ----------
+    model : tf.keras.Model
+        The user model to prune.
+    config_list : list of JSON object
+        User configuration. The format is detailed in tutorial.
+    """
     def __init__(self, model, config_list):
         super().__init__(PrunerLayerWrapper, model, config_list)
         #self.callback = PrunerCallback(self)
 
     def compress(self):
-        self.update_mask()
+        """
+        Apply compression on a pre-trained model.
+
+        If you want to prune the model during training, use callback API instead.
+
+        Returns
+        -------
+        tf.keras.Model
+            The compressed model, for convenience. This is exactly the same object to constructor argument.
+        """
+        self._update_mask()
         return self.bound_model
 
-    def update_mask(self):
+    def calc_masks(self, wrapper, **kwargs):
+        """
+        Abstract method to be overridden by algorithm. End users should ignore it.
+
+        If the callback is set up, this method will be invoked at end of each training minibatch.
+        If not, it will only be called when end user invokes ``compress``.
+
+        Parameters
+        ----------
+        wrapper : PrunerLayerWrapper
+            The instrumented layer.
+        **kwargs
+            Reserved for forward compatibility.
+
+        Returns
+        -------
+        dict of (str, tf.Tensor), or None
+            The key is weight ``Variable``'s name. The value is a mask ``Tensor`` of weight's shape and dtype.
+            If a weight's key does not appear in the return value, that weight will not be pruned.
+            Returning ``None`` means the mask is not changed since last time.
+            Weight names are globally unique, e.g. `model/conv_1/kernel:0`.
+        """
+        # TODO: maybe it should be able to calc on weight-granularity, beside from layer-granularity
+        raise NotImplementedError("Pruners must overload calc_masks()")
+
+    def _update_mask(self):
         for wrapper_idx, wrapper in enumerate(self.wrappers):
             masks = self.calc_masks(wrapper, wrapper_idx=wrapper_idx)
             if masks is not None:
                 wrapper.masks = masks
 
-    def calc_masks(self, wrapper, **kwargs):
-        # TODO: maybe it should be able to calc on weight-granularity, beside from layer-granularity
-        raise NotImplementedError("Pruners must overload calc_masks()")
-
 
 class PrunerLayerWrapper(tf.keras.Model):
+    """
+    Instrumented TF layer.
+
+    Wrappers will be passed to pruner's ``calc_masks`` API,
+    and the pruning algorithm should use wrapper's attributes to calculate masks.
+
+    Once instrumented, underlying layer's weights will get **modified** by masks before forward pass.
+
+    Attributes
+    ----------
+    layer_info : LayerInfo
+        All static information of the original layer.
+    layer : tf.keras.layers.Layer
+        The original layer.
+    config : JSON object
+        Selected configuration. The format is detailed in tutorial.
+    pruner : Pruner
+        Bound pruner object.
+    masks : dict of (str, tf.Tensor)
+        Current masks. The key is weight's name and the value is mask tensor.
+        On initialization, `masks` is an empty dict, which means no weight is pruned.
+        Afterwards, `masks` is the last return value of ``Pruner.calc_masks``.
+        See ``Pruner.calc_masks`` for details.
+    """
     def __init__(self, layer_info, pruner):
         super().__init__()
         self.layer_info = layer_info
@@ -90,6 +222,7 @@ def call(self, *inputs):
 
 
 def _detect_layers_to_compress(model, config_list):
+    # Returns list of LayerInfo.
     located_layers = _locate_layers(model)
     ret = []
     for layer in model.layers:
@@ -105,7 +238,12 @@ def _detect_layers_to_compress(model, config_list):
     return ret
 
 def _locate_layers(model, cur_path=[]):
-    # FIXME: this cannot find layers contained in list, dict, non-model custom classes, etc
+    # Find out how to access layers from model object.
+    # Returns dict of (layer's object ID, LayerInfo).
+    # This function is required because TF framework does not track layer's attribute name,
+    # 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.
     ret = {}
 
     if isinstance(model, tf.keras.Model):
@@ -129,6 +267,8 @@ def _locate_layers(model, cur_path=[]):
     return ret
 
 def _select_config(layer_info, config_list):
+    # Find the last matching config block for given layer.
+    # Returns None if the layer should not be compressed.
     ret = None
     for config in config_list:
         if 'op_types' in config:
@@ -145,7 +285,8 @@ def _select_config(layer_info, config_list):
 
 
 def _instrument_model(model, wrappers):
-    for wrapper in wrappers:
+    # Replace layers to wrappers
+    for wrapper in reversed(wrappers):
         cur = model
         for key in wrapper.layer_info.path[:-1]:
             if isinstance(key, int):
diff --git a/src/sdk/pynni/tests/test_compressor.py b/src/sdk/pynni/tests/test_compressor_torch.py
similarity index 100%
rename from src/sdk/pynni/tests/test_compressor.py
rename to src/sdk/pynni/tests/test_compressor_torch.py

From d51a6c88b989808c3aba11c66e91fcf30b803f6e Mon Sep 17 00:00:00 2001
From: liuzhe <zhliu1@microsoft.com>
Date: Wed, 12 Aug 2020 09:45:21 +0800
Subject: [PATCH 7/9] update doc

---
 src/sdk/pynni/nni/compression/tensorflow/compressor.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/sdk/pynni/nni/compression/tensorflow/compressor.py b/src/sdk/pynni/nni/compression/tensorflow/compressor.py
index 6d9679e82d..bbe4a21a52 100644
--- a/src/sdk/pynni/nni/compression/tensorflow/compressor.py
+++ b/src/sdk/pynni/nni/compression/tensorflow/compressor.py
@@ -52,7 +52,7 @@ class Compressor:
     Common base class for all compressors.
 
     This class is designed for other base classes.
-    Algorithms should inherit ``Pruner`` or Quantizer instead.
+    Algorithms should inherit ``Pruner`` or ``Quantizer`` instead.
 
 
     Attributes
@@ -99,7 +99,7 @@ class Pruner(Compressor):
     """
     Base class for pruning algorithms.
 
-    End users should use ``compress`` and callback APIs to prune their models.
+    End users should use ``compress`` and callback APIs (WIP) to prune their models.
 
     The underlying model is instrumented upon initialization of pruner object.
     So if you want to pre-train the model, train it before creating pruner object.
@@ -123,7 +123,7 @@ def compress(self):
         """
         Apply compression on a pre-trained model.
 
-        If you want to prune the model during training, use callback API instead.
+        If you want to prune the model during training, use callback API (WIP) instead.
 
         Returns
         -------

From 9ea5503bfe06d90937978b1a03960436d9ea5be4 Mon Sep 17 00:00:00 2001
From: liuzhe <zhliu1@microsoft.com>
Date: Wed, 12 Aug 2020 10:37:17 +0800
Subject: [PATCH 8/9] fix doc

---
 docs/en_US/Compressor/Pruner.md | 26 ++------------------------
 docs/zh_CN/Compressor/Pruner.md | 12 +-----------
 2 files changed, 3 insertions(+), 35 deletions(-)

diff --git a/docs/en_US/Compressor/Pruner.md b/docs/en_US/Compressor/Pruner.md
index e059834eca..9efcce8e7b 100644
--- a/docs/en_US/Compressor/Pruner.md
+++ b/docs/en_US/Compressor/Pruner.md
@@ -38,7 +38,7 @@ Tensorflow code
 ```python
 from nni.compression.tensorflow import LevelPruner
 config_list = [{ 'sparsity': 0.8, 'op_types': ['default'] }]
-pruner = LevelPruner(model_graph, config_list)
+pruner = LevelPruner(model, config_list)
 pruner.compress()
 ```
 
@@ -117,17 +117,6 @@ FPGMPruner prune filters with the smallest geometric median.
 
 ### Usage
 
-Tensorflow code
-```python
-from nni.compression.tensorflow import FPGMPruner
-config_list = [{
-    'sparsity': 0.5,
-    'op_types': ['Conv2D']
-}]
-pruner = FPGMPruner(model, config_list)
-pruner.compress()
-```
-
 PyTorch code
 ```python
 from nni.compression.torch import FPGMPruner
@@ -146,11 +135,6 @@ pruner.compress()
 ..  autoclass:: nni.compression.torch.FPGMPruner
 ```
 
-##### Tensorflow
-```eval_rst
-..  autoclass:: nni.compression.tensorflow.FPGMPruner
-```
-
 ## L1Filter Pruner
 
 This is an one-shot pruner, In ['PRUNING FILTERS FOR EFFICIENT CONVNETS'](https://arxiv.org/abs/1608.08710), authors Hao Li, Asim Kadav, Igor Durdanovic, Hanan Samet and Hans Peter Graf.
@@ -383,12 +367,6 @@ You can view [example](https://github.com/microsoft/nni/blob/master/examples/mod
 ..  autoclass:: nni.compression.torch.AGPPruner
 ```
 
-##### Tensorflow
-
-```eval_rst
-..  autoclass:: nni.compression.tensorflow.AGPPruner
-```
-
 ***
 
 ## NetAdapt Pruner
@@ -620,4 +598,4 @@ pruner.compress(eval_args=[model], finetune_args=[model])
 
 ```eval_rst
 ..  autoclass:: nni.compression.torch.SensitivityPruner
-```
\ No newline at end of file
+```
diff --git a/docs/zh_CN/Compressor/Pruner.md b/docs/zh_CN/Compressor/Pruner.md
index f78b8e0f1c..d11829e4b5 100644
--- a/docs/zh_CN/Compressor/Pruner.md
+++ b/docs/zh_CN/Compressor/Pruner.md
@@ -37,7 +37,7 @@ TensorFlow 代码
 ```python
 from nni.compression.tensorflow import LevelPruner
 config_list = [{ 'sparsity': 0.8, 'op_types': ['default'] }]
-pruner = LevelPruner(model_graph, config_list)
+pruner = LevelPruner(model, config_list)
 pruner.compress()
 ```
 
@@ -102,16 +102,6 @@ pruner.compress()
 
 ### 用法
 
-TensorFlow 代码
-```python
-from nni.compression.tensorflow import FPGMPruner
-config_list = [{
-    'sparsity': 0.5,
-    'op_types': ['Conv2D']
-}]
-pruner = FPGMPruner(model, config_list)
-pruner.compress()
-```
 PyTorch 代码
 ```python
 from nni.compression.torch import FPGMPruner

From 7a22fd9ef531daa32583b1546a13215931ad6b6e Mon Sep 17 00:00:00 2001
From: liuzhe <zhliu1@microsoft.com>
Date: Wed, 12 Aug 2020 10:38:26 +0800
Subject: [PATCH 9/9] fix lint

---
 src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py b/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py
index feefe95f7d..ace3d39e4e 100644
--- a/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py
+++ b/src/sdk/pynni/nni/compression/tensorflow/pruning/one_shot.py
@@ -43,7 +43,7 @@ def calc_masks(self, sparsity, wrapper, wrapper_idx=None):
 class LevelPrunerMasker(WeightMasker):
     def calc_masks(self, sparsity, wrapper, wrapper_idx=None):
         masks = {}
-        for i, weight_variable in enumerate(wrapper.layer.weights):
+        for weight_variable in wrapper.layer.weights:
             if weight_variable.name == 'bias':
                 continue