Merge branch 'v2.0' into master

docs/requirements.txt

-sphinx==3.3.1
-sphinx-argparse==0.2.5
-sphinx-rtd-theme==0.4.2
-sphinxcontrib-websupport==1.1.0
-pygments==2.7.1
+sphinx>=3.3.1
+sphinx-argparse
+sphinx-rtd-theme
+sphinxcontrib-websupport
+pygments>=2.7.1
 hyperopt
 json_tricks
 numpy
@@ -14,4 +14,12 @@ schema
 tensorboard
 scikit-learn>=0.23.2
 thop
-https://download.pytorch.org/whl/cpu/torch-1.3.1%2Bcpu-cp37-cp37m-linux_x86_64.whl
+colorama
+pkginfo
+websockets
+filelock
+prettytable
+psutil
+ruamel.yaml
+https://download.pytorch.org/whl/cpu/torch-1.7.1%2Bcpu-cp37-cp37m-linux_x86_64.whl
+https://download.pytorch.org/whl/cpu/torchvision-0.8.2%2Bcpu-cp37-cp37m-linux_x86_64.whl

examples/feature_engineering/gbdt_selector/gbdt_selector_test.py

@@ -22,7 +22,7 @@ import numpy as np
 from sklearn.datasets import load_svmlight_file
 from sklearn.model_selection import train_test_split
 
-from nni.feature_engineering.gbdt_selector import GBDTSelector
+from nni.algorithms.feature_engineering.gbdt_selector import GBDTSelector
 
 url_zip_train = 'https://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/binary/rcv1_train.binary.bz2'
 urllib.request.urlretrieve(url_zip_train, filename='train.bz2')

examples/feature_engineering/gradient_feature_selector/benchmark_test.py

@@ -33,7 +33,7 @@ from sklearn.linear_model import LogisticRegression
 from sklearn.ensemble import ExtraTreesClassifier
 from sklearn.feature_selection import SelectFromModel
 
-from nni.feature_engineering.gradient_selector import FeatureGradientSelector
+from nni.algorithms.feature_engineering.gradient_selector import FeatureGradientSelector
 
 
 class Benchmark():

examples/feature_engineering/gradient_feature_selector/sklearn_test.py

@@ -28,7 +28,7 @@ from sklearn.linear_model import LogisticRegression
 from sklearn.ensemble import ExtraTreesClassifier
 from sklearn.feature_selection import SelectFromModel
 
-from nni.feature_engineering.gradient_selector import FeatureGradientSelector
+from nni.algorithms.feature_engineering.gradient_selector import FeatureGradientSelector
 
 
 def test():

examples/model_compress/fpgm_tf_mnist.py

-import tensorflow as tf
-from tensorflow import keras
-assert tf.__version__ >= "2.0"
-import numpy as np
-from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
-from nni.compression.tensorflow import FPGMPruner
-
-def get_data():
-    (X_train_full, y_train_full), _ = keras.datasets.mnist.load_data()
-    X_train, X_valid = X_train_full[:-5000], X_train_full[-5000:]
-    y_train, y_valid = y_train_full[:-5000], y_train_full[-5000:]
-
-    X_mean = X_train.mean(axis=0, keepdims=True)
-    X_std = X_train.std(axis=0, keepdims=True) + 1e-7
-    X_train = (X_train - X_mean) / X_std
-    X_valid = (X_valid - X_mean) / X_std
-
-    X_train = X_train[..., np.newaxis]
-    X_valid = X_valid[..., np.newaxis]
-
-    return X_train, X_valid, y_train, y_valid
-
-def get_model():
-    model = keras.models.Sequential([
-        Conv2D(filters=32, kernel_size=7, input_shape=[28, 28, 1], activation='relu', padding="SAME"),
-        MaxPooling2D(pool_size=2),
-        Conv2D(filters=64, kernel_size=3, activation='relu', padding="SAME"),
-        MaxPooling2D(pool_size=2),
-        Flatten(),
-        Dense(units=128, activation='relu'),
-        Dropout(0.5),
-        Dense(units=10, activation='softmax'),
-    ])
-    model.compile(loss="sparse_categorical_crossentropy",
-        optimizer=keras.optimizers.SGD(lr=1e-3),
-        metrics=["accuracy"])
-    return model
-
-def main():
-    X_train, X_valid, y_train, y_valid = get_data()
-    model = get_model()
-
-    configure_list = [{
-        'sparsity': 0.5,
-        'op_types': ['Conv2D']
-    }]
-    pruner = FPGMPruner(model, configure_list)
-    pruner.compress()
-
-    update_epoch_callback = keras.callbacks.LambdaCallback(on_epoch_begin=lambda epoch, logs: pruner.update_epoch(epoch))
-
-    model.fit(X_train, y_train, epochs=10, validation_data=(X_valid, y_valid), callbacks=[update_epoch_callback])
-
-
-if __name__ == '__main__':
-    main()

examples/model_compress/knowledge_distill/knowledge_distill.py

-import logging
-import torch
-import torch.nn.functional as F
-
-_logger = logging.getLogger(__name__)
-
-
-class KnowledgeDistill():
-    """
-    Knowledge Distillaion support while fine-tuning the compressed model
-    Geoffrey Hinton, Oriol Vinyals, Jeff Dean
-    "Distilling the Knowledge in a Neural Network"
-    https://arxiv.org/abs/1503.02531
-    """
-
-    def __init__(self, teacher_model, kd_T=1):
-        """
-        Parameters
-        ----------
-        teacher_model : pytorch model
-            the teacher_model for teaching the student model, it should be pretrained
-        kd_T: float
-            kd_T is the temperature parameter, when kd_T=1 we get the standard softmax function
-            As kd_T grows, the probability distribution generated by the softmax function becomes softer
-        """
-
-        self.teacher_model = teacher_model
-        self.kd_T = kd_T
-
-    def _get_kd_loss(self, data, student_out, teacher_out_preprocess=None):
-        """
-        Parameters
-        ----------
-        data : torch.Tensor
-            the input training data
-        student_out: torch.Tensor
-            output of the student network
-        teacher_out_preprocess: function
-            a function for pre-processing teacher_model's output
-            e.g. when teacher_out_preprocess=lambda x:x[0]
-            extract teacher_model's output (tensor1, tensor2)->tensor1
-
-        Returns
-        -------
-        torch.Tensor
-            weighted distillation loss
-        """
-
-        with torch.no_grad():
-            kd_out = self.teacher_model(data)
-        if teacher_out_preprocess is not None:
-            kd_out = teacher_out_preprocess(kd_out)
-        assert type(kd_out) is torch.Tensor
-        assert type(student_out) is torch.Tensor
-        assert kd_out.shape == student_out.shape
-        soft_log_out = F.log_softmax(student_out / self.kd_T, dim=1)
-        soft_t = F.softmax(kd_out / self.kd_T, dim=1)
-        loss_kd = F.kl_div(soft_log_out, soft_t.detach(), reduction='batchmean')
-        return loss_kd
-
-    def loss(self, data, student_out):
-        """
-        Parameters
-        ----------
-        data : torch.Tensor
-            Input of the student model
-        student_out : torch.Tensor
-            Output of the student model
-
-        Returns
-        -------
-        torch.Tensor
-            Weighted loss of student loss and distillation loss
-        """
-        return self._get_kd_loss(data, student_out)

examples/model_compress/model_prune_tf.py

-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.9,
-            '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),
-        tf.keras.layers.Softmax()
-    ])
-
-
-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)
-
-
-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)
-
-    print('start training')
-    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'],
-        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__':
-    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)