merge master

.gitignore

@@ -11,6 +11,9 @@ pids
 *.seed
 *.pid.lock
 
+# Build package
+dist/
+
 # Directory for instrumented libs generated by jscoverage/JSCover
 lib-cov
 
@@ -54,9 +57,6 @@ typings/
 # Yarn Integrity file
 .yarn-integrity
 
-# dotenv environment variables file
-.env
-
 # next.js build output
 .next
 
@@ -67,7 +67,18 @@ typings/
 __pycache__
 build
 *.egg-info
+setup.pye
 
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# VSCode
 .vscode
 
 # In case you place source code in ~/nni/

Makefile

@@ -52,6 +52,7 @@ $(shell mkdir -p $(NNI_DEPENDENCY_FOLDER))
 NNI_NODE_TARBALL ?= $(NNI_DEPENDENCY_FOLDER)/nni-node-$(OS_SPEC)-x64.tar.xz
 NNI_NODE_FOLDER = $(NNI_DEPENDENCY_FOLDER)/nni-node-$(OS_SPEC)-x64
 NNI_NODE ?= $(BIN_FOLDER)/node
+NNI_NPM ?= $(BIN_FOLDER)/npm
 NNI_YARN_TARBALL ?= $(NNI_DEPENDENCY_FOLDER)/nni-yarn.tar.gz
 NNI_YARN_FOLDER ?= $(NNI_DEPENDENCY_FOLDER)/nni-yarn
 NNI_YARN ?= PATH=$(BIN_FOLDER):$${PATH} $(NNI_YARN_FOLDER)/bin/yarn
@@ -149,8 +150,9 @@ install-dependencies: $(NNI_NODE_TARBALL) $(NNI_YARN_TARBALL)
 	mkdir $(NNI_NODE_FOLDER)
 	tar -xf $(NNI_NODE_TARBALL) -C $(NNI_NODE_FOLDER) --strip-components 1
 	mkdir -p $(BIN_FOLDER)
-	rm -f $(NNI_NODE)
-	cp $(NNI_NODE_FOLDER)/bin/node $(NNI_NODE)
+	rm -f $(NNI_NODE) $(NNI_NPM)
+	ln -s $(NNI_NODE_FOLDER)/bin/node $(NNI_NODE)
+	ln -s $(NNI_NODE_FOLDER)/bin/npm $(NNI_NPM)
 
 	#$(_INFO) Extracting Yarn $(_END)
 	rm -rf $(NNI_YARN_FOLDER)

azure-pipelines.yml

@@ -8,16 +8,34 @@ jobs:
         PYTHON_VERSION: '3.6'
 
   steps:
-  - script: python3 -m pip install --upgrade pip setuptools --user
+  - script: |
+      python3 -m pip install --upgrade pip setuptools --user
+      python3 -m pip install pylint==2.3.1 astroid==2.2.5 --user
+      python3 -m pip install coverage --user
+      echo "##vso[task.setvariable variable=PATH]${HOME}/.local/bin:${PATH}"
     displayName: 'Install python tools'
+  - script: |
+      source install.sh
+    displayName: 'Install nni toolkit via source code'
   - script: |
       python3 -m pip install torch==0.4.1 --user
       python3 -m pip install torchvision==0.2.1 --user
       python3 -m pip install tensorflow==1.13.1 --user
+      python3 -m pip install keras==2.1.6 --user
+      python3 -m pip install gym onnx --user
+      sudo apt-get install swig -y
+      nnictl package install --name=SMAC
+      nnictl package install --name=BOHB
     displayName: 'Install dependencies'
   - script: |
-      source install.sh
-    displayName: 'Install nni toolkit via source code'
+      set -e
+      python3 -m pylint --rcfile pylintrc nni_annotation
+      python3 -m pylint --rcfile pylintrc nni_cmd
+      python3 -m pylint --rcfile pylintrc nni_gpu_tool
+      python3 -m pylint --rcfile pylintrc nni_trial_tool
+      python3 -m pylint --rcfile pylintrc nni
+      python3 -m pylint --rcfile pylintrc nnicli
+    displayName: 'Run pylint'
   - script: |
       python3 -m pip install flake8 --user
       IGNORE=./tools/nni_annotation/testcase/*:F821,./examples/trials/mnist-nas/*/mnist*.py:F821,./examples/trials/nas_cifar10/src/cifar10/general_child.py:F821
@@ -29,19 +47,19 @@ jobs:
     displayName: 'Unit test'
   - script: |
       cd test
-      PATH=$HOME/.local/bin:$PATH python3 naive_test.py
+      python3 naive_test.py
     displayName: 'Naive test'
   - script: |
       cd test
-      PATH=$HOME/.local/bin:$PATH python3 tuner_test.py
+      python3 tuner_test.py
     displayName: 'Built-in tuners / assessors tests'
   - script: |
       cd test
-      PATH=$HOME/.local/bin:$PATH python3 metrics_test.py
+      python3 metrics_test.py
     displayName: 'Trial job metrics test'
   - script: |
       cd test
-      PATH=$HOME/.local/bin:$PATH python3 cli_test.py
+      python3 cli_test.py
     displayName: 'nnicli test'
 
 - job: 'basic_test_pr_macOS'
@@ -55,29 +73,34 @@ jobs:
   steps:
   - script: python3 -m pip install --upgrade pip setuptools
     displayName: 'Install python tools'
+  - script: |
+      source install.sh
+      echo "##vso[task.setvariable variable=PATH]${HOME}/Library/Python/3.7/bin:${PATH}"
+    displayName: 'Install nni toolkit via source code'
   - script: |
       python3 -m pip install torch==0.4.1 --user
       python3 -m pip install torchvision==0.2.1 --user
       python3 -m pip install tensorflow==1.13.1 --user
+      ruby -e "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/master/install)" < /dev/null 2> /dev/null
+      brew install swig@3
+      ln -s /usr/local/opt/swig\@3/bin/swig /usr/local/bin/swig
+      nnictl package install --name=SMAC
     displayName: 'Install dependencies'
-  - script: |
-      source install.sh
-    displayName: 'Install nni toolkit via source code'
   - script: |
       cd test
-      PATH=$HOME/Library/Python/3.7/bin:$PATH && source unittest.sh
+      source unittest.sh
     displayName: 'Unit test'
   - script: |
       cd test
-      PATH=$HOME/Library/Python/3.7/bin:$PATH python3 naive_test.py
+      python3 naive_test.py
     displayName: 'Naive test'
   - script: |
       cd test
-      PATH=$HOME/Library/Python/3.7/bin:$PATH python3 tuner_test.py
+      python3 tuner_test.py
     displayName: 'Built-in tuners / assessors tests'
   - script: |
       cd test
-      PATH=$HOME/Library/Python/3.7/bin:$PATH python3 cli_test.py
+      python3 cli_test.py
     displayName: 'nnicli test'
 
 - job: 'basic_test_pr_Windows'

deployment/pypi/.gitignore

@@ -1,6 +1,8 @@
 nni/
+nni-yarn/
 dist/
 build/
 *.egg-info/
+nni-yarn.tar.gz
 node-*.tar.xz
-node-*/
+node-*/

docs/en_US/Compressor/AutoCompression.md

@@ -9,13 +9,13 @@ You can easily compress a model with NNI compression. Take pruning for example,
 ```python
 from nni.compression.torch import LevelPruner
 config_list = [{ 'sparsity': 0.8, 'op_types': ['default'] }]
-pruner = LevelPruner(config_list)
-pruner(model)
+pruner = LevelPruner(model, config_list)
+pruner.compress()
 ```
 
 The 'default' op_type stands for the module types defined in [default_layers.py](https://github.com/microsoft/nni/blob/master/src/sdk/pynni/nni/compression/torch/default_layers.py) for pytorch.
 
-Therefore ```{ 'sparsity': 0.8, 'op_types': ['default'] }```means that **all layers with specified op_types will be compressed with the same 0.8 sparsity**. When ```pruner(model)``` called, the model is compressed with masks and after that you can normally fine tune this model and **pruned weights won't be updated** which have been masked.
+Therefore ```{ 'sparsity': 0.8, 'op_types': ['default'] }```means that **all layers with specified op_types will be compressed with the same 0.8 sparsity**. When ```pruner.compress()``` called, the model is compressed with masks and after that you can normally fine tune this model and **pruned weights won't be updated** which have been masked.
 
 ## Then, make this automatic
 
@@ -84,9 +84,9 @@ config_list_agp = [{'initial_sparsity': 0, 'final_sparsity': conv0_sparsity,
                    {'initial_sparsity': 0, 'final_sparsity': conv1_sparsity,
                     'start_epoch': 0, 'end_epoch': 3,
                     'frequency': 1,'op_name': 'conv1' },]
-PRUNERS = {'level':LevelPruner(config_list_level)，'agp':AGP_Pruner(config_list_agp)}
+PRUNERS = {'level':LevelPruner(model, config_list_level)，'agp':AGP_Pruner(model, config_list_agp)}
 pruner = PRUNERS(params['prune_method']['_name'])
-pruner(model)
+pruner.compress()
 ... # fine tuning
 acc = evaluate(model) # evaluation
 nni.report_final_results(acc)

docs/en_US/Compressor/Overview.md

@@ -25,22 +25,22 @@ Tensorflow code
 ```python
 from nni.compression.tensorflow import LevelPruner
 config_list = [{ 'sparsity': 0.8, 'op_types': ['default'] }]
-pruner = LevelPruner(config_list)
-pruner(tf.get_default_graph())
+pruner = LevelPruner(tf.get_default_graph(), config_list)
+pruner.compress()
 ```
 
 PyTorch code
 
 ```python
 from nni.compression.torch import LevelPruner
 config_list = [{ 'sparsity': 0.8, 'op_types': ['default'] }]
-pruner = LevelPruner(config_list)
-pruner(model)
+pruner = LevelPruner(model, config_list)
+pruner.compress()
 ```
 
 You can use other compression algorithms in the package of `nni.compression`. The algorithms are implemented in both PyTorch and Tensorflow, under `nni.compression.torch` and `nni.compression.tensorflow` respectively. You can refer to [Pruner](./Pruner.md) and [Quantizer](./Quantizer.md) for detail description of supported algorithms.
 
-The function call `pruner(model)` receives user defined model (in Tensorflow the model can be obtained with `tf.get_default_graph()`, while in PyTorch the model is the defined model class), and the model is modified with masks inserted. Then when you run the model, the masks take effect. The masks can be adjusted at runtime by the algorithms.
+The function call `pruner.compress()` modifies user defined model (in Tensorflow the model can be obtained with `tf.get_default_graph()`, while in PyTorch the model is the defined model class), and the model is modified with masks inserted. Then when you run the model, the masks take effect. The masks can be adjusted at runtime by the algorithms.
 
 When instantiate a compression algorithm, there is `config_list` passed in. We describe how to write this config below.
 
@@ -95,7 +95,17 @@ pruner.update_epoch(epoch)
 
 The other is `step`, it can be called with `pruner.step()` after each minibatch. Note that not all algorithms need these two APIs, for those that do not need them, calling them is allowed but has no effect.
 
-__[TODO]__ The last API is for users to export the compressed model. You will get a compressed model when you finish the training using this API. It also exports another file storing the values of masks.
+You can easily export the compressed model using the following API if you are pruning your model, ```state_dict``` of the sparse model weights will be stored in ```model.pth```, which can be loaded by ```torch.load('model.pth')```
+
+```
+pruner.export_model(model_path='model.pth')
+```
+
+```mask_dict ``` and pruned model in ```onnx``` format(```input_shape``` need to be specified) can also be exported like this:
+
+```python
+pruner.export_model(model_path='model.pth', mask_path='mask.pth', onnx_path='model.onnx', input_shape=[1, 1, 28, 28])
+```
 
 ## Customize new compression algorithms
 
@@ -111,20 +121,26 @@ If you want to write a new pruning algorithm, you can write a class that inherit
 # nni.compression.tensorflow.Pruner with
 # nni.compression.torch.Pruner
 class YourPruner(nni.compression.tensorflow.Pruner):
-    def __init__(self, config_list):
-        # suggest you to use the NNI defined spec for config
-        super().__init__(config_list)
-
-    def bind_model(self, model):
-        # this func can be used to remember the model or its weights
-        # in member variables, for getting their values during training
-        pass
-
-    def calc_mask(self, weight, config, **kwargs):
-        # weight is the target weight tensor
-        # config is the selected dict object in config_list for this layer
-        # kwargs contains op, op_types, and op_name
-        # design your mask and return your mask
+    def __init__(self, model, config_list):
+        """
+        Suggest you to use the NNI defined spec for config
+        """
+        super().__init__(model, config_list)
+
+    def calc_mask(self, layer, config):
+        """
+        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
+        """
         return your_mask
 
     # note for pytorch version, there is no sess in input arguments
@@ -133,16 +149,18 @@ class YourPruner(nni.compression.tensorflow.Pruner):
 
     # note for pytorch version, there is no sess in input arguments
     def step(self, sess):
-        # can do some processing based on the model or weights binded
-        # in the func bind_model
+        """
+        Can do some processing based on the model or weights binded
+        in the func bind_model
+        """
         pass
 ```
 
-For the simplest algorithm, you only need to override `calc_mask`. It receives each layer's weight and selected configuration, as well as op information. You generate the mask for this weight in this function and return. Then NNI applies the mask for you.
+For the simplest algorithm, you only need to override ``calc_mask``. It receives the to-be-compressed layers one by one along with their compression configuration. You generate the mask for this weight in this function and return. Then NNI applies the mask for you.
 
-Some algorithms generate mask based on training progress, i.e., epoch number. We provide `update_epoch` for the pruner to be aware of the training progress.
+Some algorithms generate mask based on training progress, i.e., epoch number. We provide `update_epoch` for the pruner to be aware of the training progress. It should be called at the beginning of each epoch.
 
-Some algorithms may want global information for generating masks, for example, all weights of the model (for statistic information), model optimizer's information. NNI supports this requirement using `bind_model`. `bind_model` receives the complete model, thus, it could record any information (e.g., reference to weights) it cares about. Then `step` can process or update the information according to the algorithm. You can refer to [source code of built-in algorithms](https://github.com/microsoft/nni/tree/master/src/sdk/pynni/nni/compressors) for example implementations.
+Some algorithms may want global information for generating masks, for example, all weights of the model (for statistic information). Your can use `self.bound_model` in the Pruner class for accessing weights. If you also need optimizer's information (for example in Pytorch), you could override `__init__` to receive more arguments such as model's optimizer. Then `step` can process or update the information according to the algorithm. You can refer to [source code of built-in algorithms](https://github.com/microsoft/nni/tree/master/src/sdk/pynni/nni/compressors) for example implementations.
 
 ### Quantization algorithm
 
@@ -154,20 +172,19 @@ The interface for customizing quantization algorithm is similar to that of pruni
 # nni.compression.tensorflow.Quantizer with
 # nni.compression.torch.Quantizer
 class YourQuantizer(nni.compression.tensorflow.Quantizer):
-    def __init__(self, config_list):
-        # suggest you to use the NNI defined spec for config
-        super().__init__(config_list)
-
-    def bind_model(self, model):
-        # this func can be used to remember the model or its weights
-        # in member variables, for getting their values during training
-        pass
+    def __init__(self, model, config_list):
+        """
+        Suggest you to use the NNI defined spec for config
+        """
+        super().__init__(model, config_list)
 
     def quantize_weight(self, weight, config, **kwargs):
-        # weight is the target weight tensor
-        # config is the selected dict object in config_list for this layer
-        # kwargs contains op, op_types, and op_name
-        # design your quantizer and return new weight
+        """
+        weight is the target weight tensor
+        config is the selected dict object in config_list for this layer
+        kwargs contains op, op_types, and op_name
+        design your quantizer and return new weight
+        """
         return new_weight
 
     # note for pytorch version, there is no sess in input arguments
@@ -176,8 +193,10 @@ class YourQuantizer(nni.compression.tensorflow.Quantizer):
 
     # note for pytorch version, there is no sess in input arguments
     def step(self, sess):
-        # can do some processing based on the model or weights binded
-        # in the func bind_model
+        """
+        Can do some processing based on the model or weights binded
+        in the func bind_model
+        """
         pass
 ```