Analysis utils

deployment/pypi/setup.py

@@ -62,7 +62,9 @@
         'scipy',
         'coverage',
         'colorama',
-        'scikit-learn>=0.20,<0.22'
+        'scikit-learn>=0.20,<0.22',
+        'graphviz',
+        'matplotlib'
     ],
     classifiers = [
         'Programming Language :: Python :: 3',

setup.py

@@ -40,7 +40,9 @@ def read(fname):
         'schema',
         'PythonWebHDFS',
         'colorama',
-        'scikit-learn>=0.20,<0.22'
+        'scikit-learn>=0.20,<0.22',
+        'graphviz',
+        'matplotlib'
     ],
 
     entry_points = {

src/sdk/pynni/nni/analysis_utils/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.

src/sdk/pynni/nni/analysis_utils/sensitivity/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.

src/sdk/pynni/nni/analysis_utils/sensitivity/torch/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+from .sensitivity_analysis import SensitivityAnalysis

src/sdk/pynni/nni/analysis_utils/sensitivity/torch/sensitivity_analysis.py

@@ -0,0 +1,257 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import os
+import copy
+import csv
+import logging
+from collections import OrderedDict
+import matplotlib
+import matplotlib.pyplot as plt
+from matplotlib.lines import Line2D
+
+import numpy as np
+import torch.nn as nn
+
+from nni.compression.torch import LevelPruner
+from nni.compression.torch import L1FilterPruner
+from nni.compression.torch import L2FilterPruner
+
+# use Agg backend
+matplotlib.use('Agg')
+SUPPORTED_OP_NAME = ['Conv2d', 'Conv1d']
+SUPPORTED_OP_TYPE = [getattr(nn, name) for name in SUPPORTED_OP_NAME]
+
+logger = logging.getLogger('Sensitivity_Analysis')
+logger.setLevel(logging.INFO)
+
+
+class SensitivityAnalysis:
+    def __init__(self, model, val_func, sparsities=None, prune_type='l1', early_stop=1.0):
+        """
+        Perform sensitivity analysis for this model.
+        Parameters
+        ----------
+            model:
+                the model to perform sensitivity analysis
+            val_func:
+                validation function for the model. Due to
+                different models may need different dataset/criterion
+                , therefore the user need to cover this part by themselves.
+                val_func take the model as the first input parameter, and
+                return the accuracy as output.
+            sparsities:
+                The sparsity list provided by users.
+            prune_type:
+                The pruner type used to prune the conv layers, default is 'l1',
+                and 'l2', 'fine-grained' is also supported.
+            early_stop:
+                If this flag is set, the sensitivity analysis
+                for a conv layer will early stop when the accuracy
+                drop already reach the value of early_stop (0.05 for example).
+                The default value is 1.0, which means the analysis won't stop
+                until all given sparsities are tested.
+
+        """
+        self.model = model
+        self.val_func = val_func
+        self.target_layer = OrderedDict()
+        self.ori_state_dict = copy.deepcopy(self.model.state_dict())
+        self.target_layer = {}
+        self.sensitivities = {}
+        if sparsities is not None:
+            self.sparsities = sorted(sparsities)
+        else:
+            self.sparsities = np.arange(0.1, 1.0, 0.1)
+        self.sparsities = [np.round(x, 2) for x in self.sparsities]
+        self.Pruner = L1FilterPruner
+        if prune_type == 'l2':
+            self.Pruner = L2FilterPruner
+        elif prune_type == 'fine-grained':
+            self.Pruner = LevelPruner
+        self.early_stop = early_stop
+        self.ori_acc = None  # original accuracy for the model
+        # already_pruned is for the iterative sensitivity analysis
+        # For example, sensitivity_pruner iteratively prune the target
+        # model according to the sensitivity. After each round of
+        # pruning, the sensitivity_pruner will test the new sensitivity
+        # for each layer
+        self.already_pruned = {}
+        self.model_parse()
+
+    @property
+    def layers_count(self):
+        return len(self.target_layer)
+
+    def model_parse(self):
+        for name, submodel in self.model.named_modules():
+            for op_type in SUPPORTED_OP_TYPE:
+                if isinstance(submodel, op_type):
+                    self.target_layer[name] = submodel
+                    self.already_pruned[name] = 0
+
+    def analysis(self, val_args=None, val_kwargs=None, start=0, end=None):
+        """
+        This function analyze the sensitivity to pruning for
+        each conv layer in the target model.
+        If %start and %end are not set, we analyze all the conv
+        layers by default. Users can specify several layers to
+        analyze or parallelize the analysis process easily through
+        the %start and %end parameter.
+
+        Parameters
+        ----------
+            start:
+                Layer index of the sensitivity analysis start
+            end:
+                Layer index of the sensitivity analysis end
+            val_args:
+                args for the val_function
+            val_kwargs:
+                kwargs for the val_funtion
+                The val_funtion will be called as:
+                    val_function(*val_args, **val_kwargs)
+
+        Returns
+        -------
+            sensitivities:
+                dict object that stores the trajectory of the
+                accuracy when the prune ratio changes
+        """
+        if not end:
+            end = self.layers_count
+        assert start >= 0 and end <= self.layers_count
+        assert start <= end
+        if val_args is None:
+            val_args = []
+        if val_kwargs is None:
+            val_kwargs = {}
+        # Get the validation accuracy before pruning
+        if self.ori_acc is None:
+            self.ori_acc = self.val_func(*val_args, **val_kwargs)
+        namelist = list(self.target_layer.keys())
+        for layerid in range(start, end):
+            name = namelist[layerid]
+            self.sensitivities[name] = {}
+            for sparsity in self.sparsities:
+                # Calculate the actual prune ratio based on the already pruned ratio
+                sparsity = (
+                    1.0 - self.already_pruned[name]) * sparsity + self.already_pruned[name]
+                # TODO In current L1/L2 Filter Pruner, the 'op_types' is still necessary
+                # I think the L1/L2 Pruner should specify the op_types automaticlly
+                # according to the op_names
+                cfg = [{'sparsity': sparsity, 'op_names': [
+                    name], 'op_types': ['Conv2d']}]
+                pruner = self.Pruner(self.model, cfg)
+                pruner.compress()
+                val_acc = self.val_func(*val_args, **val_kwargs)
+                logger.info('Layer: %s Sparsity: %.2f Accuracy: %.4f',
+                            name, sparsity, val_acc)
+
+                self.sensitivities[name][sparsity] = val_acc
+                pruner._unwrap_model()
+                del pruner
+                # if the accuracy drop already reach the 'early_stop'
+                if val_acc + self.early_stop < self.ori_acc:
+                    break
+
+            # reset the weights pruned by the pruner, because
+            # out sparsities is sorted, so we donnot need to reset
+            # weight of the layer when the sparsity changes, instead,
+            # we only need reset the weight when the pruning layer changes.
+            self.model.load_state_dict(self.ori_state_dict)
+
+        return self.sensitivities
+
+    def visualization(self, outdir, merge=False):
+        """
+        Visualize the sensitivity curves of the model
+
+        Parameters
+        ----------
+            outdir:
+                output directory of the image
+            merge:
+                if merge all the sensitivity curves into a
+                single image. If not, we will draw a picture
+                for each target layer of the model.
+        """
+        os.makedirs(outdir, exist_ok=True)
+        LineStyles = [':', '-.', '--', '-']
+        Markers = list(Line2D.markers.keys())
+        if not merge:
+            # Draw the sensitivity curves for each layer first
+            for name in self.sensitivities:
+                X = list(self.sensitivities[name].keys())
+                X = sorted(X)
+                Y = [self.sensitivities[name][x] for x in X]
+                if 0.00 not in X:
+                    # add the original accuracy into the figure
+                    X = [0.00] + X
+                    Y = [self.ori_acc] + Y
+                plt.figure(figsize=(8, 4))
+                plt.plot(X, Y, marker='*')
+                plt.xlabel('Prune Ratio')
+                plt.ylabel('Validation Accuracy')
+                plt.title(name)
+                plt.tight_layout()
+                filepath = os.path.join(outdir, '%s.jpg' % name)
+                plt.savefig(filepath)
+                plt.close()
+        else:
+            plt.figure()
+            styleid = 0
+            for name in self.sensitivities:
+                X = list(self.sensitivities[name].keys())
+                X = sorted(X)
+                Y = [self.sensitivities[name][x] for x in X]
+                if 0.00 not in X:
+                    # add the original accuracy into the figure
+                    X = [0.00] + X
+                    Y = [self.ori_acc] + Y
+                linestyle = LineStyles[styleid % len(LineStyles)]
+                marker = Markers[styleid % len(Markers)]
+                plt.plot(X, Y, label=name, linestyle=linestyle, marker=marker)
+                plt.xlabel('Prune Ratio')
+                plt.ylabel('Validation Accuracy')
+                plt.legend(loc='center left', bbox_to_anchor=(1.02, 0.5))
+                plt.tight_layout()
+                filepath = os.path.join(outdir, 'all.jpg')
+                plt.savefig(filepath, dpi=1000, bbox_inches='tight')
+                styleid += 1
+            plt.close()
+
+    def export(self, filepath):
+        """
+        Export the results of the sensitivity analysis
+        to a csv file.
+
+        Parameters
+        ----------
+            filepath:
+                Path of the output file
+        """
+        str_sparsities = [str(x) for x in self.sparsities]
+        header = ['layername'] + str_sparsities
+        with open(filepath, 'w') as csvf:
+            csv_w = csv.writer(csvf)
+            csv_w.writerow(header)
+            for layername in self.sensitivities:
+                row = []
+                row.append(layername)
+                for sparsity in sorted(self.sensitivities[layername].keys()):
+                    row.append(self.sensitivities[layername][sparsity])
+                csv_w.writerow(row)
+
+    def update_already_pruned(self, layername, ratio):
+        """
+        Set the already pruned ratio for the target layer.
+        """
+        self.already_pruned[layername] = ratio
+
+    def load_state_dict(self, state_dict):
+        """
+        Update the weight of the model
+        """
+        self.ori_state_dict = copy.deepcopy(state_dict)
+        self.model.load_state_dict(self.ori_state_dict)

src/sdk/pynni/nni/analysis_utils/topology/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.

src/sdk/pynni/nni/analysis_utils/topology/torch/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.