PyTorch Support 🚀 🚀 🚀 (#16)

* PyTorch Support 🚀 🚀 🚀

* Fix CI

Author: tkat0

.gitignore

@@ -1,3 +1,4 @@
+tests_results/
 
 # Created by https://www.gitignore.io/api/python,pycharm
 # Edit at https://www.gitignore.io/?templates=python,pycharm

.travis.yml

@@ -1,17 +1,19 @@
 language: python
 
 python:
-  - 3.6
+  - "3.6"
 
 cache:
   - pip
 
 install:
   - pip install -qU pip setuptools
-  - pip install -q pytest==3.6 pytest-cov coverage
+  - pip install -q pytest==3.6 pytest-cov coverage matplotlib
   - pip install -q chainer networkx chainer_computational_cost
   - pip install -q chainercv
+  - pip install -q https://download.pytorch.org/whl/cpu/torch-1.1.0-cp36-cp36m-linux_x86_64.whl
+  - pip install -q https://download.pytorch.org/whl/cpu/torchvision-0.3.0-cp36-cp36m-linux_x86_64.whl
   - pip install -q -e .
 
 script:
-  - pytest -vv tests
+  - pytest -vv --cov=chainerpruner --cov-report=html --color=auto -s --basetemp tests_results tests

README.md

@@ -1,6 +1,7 @@
 # ChainerPruner
 
 [![Build Status](https://travis-ci.org/DeNA/ChainerPruner.svg?branch=master)](https://travis-ci.org/DeNA/ChainerPruner)
+![License](https://img.shields.io/badge/license-MIT-brightgreen.svg)
 
 ChainerPruner: Channel Pruning framework for [Chainer](https://github.com/chainer/chainer)
 

chainerpruner/__init__.py

@@ -5,13 +5,33 @@
 __author_email__ = '[email protected]'
 __url__ = 'https://github.com/DeNA/ChainerPruner'
 
+try:
+    import chainer
+
+    avalable_chainer = True
+except ImportError:
+    avalable_chainer = False
+
+try:
+    import torch
+
+    avalable_pytorch = True
+except ImportError:
+    avalable_pytorch = False
+
+import os
+from chainerpruner import utils
 from chainerpruner import masks
 from chainerpruner import pruning
 from chainerpruner import rebuild
 from chainerpruner import serializers
-from chainerpruner import utils
-from chainerpruner import trace
 
-from chainerpruner.graph import Graph
-from chainerpruner.node import Node
 from chainerpruner.pruner import Pruner
+from chainerpruner.graph import Graph
+
+disable_patch = os.getenv('CHAINERPRUNER_DISABLE_PYTORCH_LOAD_PATCH', False)
+
+if avalable_pytorch and not disable_patch:
+    from chainerpruner.serializers.pytorch import enable_custom_load_state_dict
+
+    enable_custom_load_state_dict()

chainerpruner/graph/__init__.py

@@ -0,0 +1 @@
+from chainerpruner.graph.graph import Graph

chainerpruner/rebuild/links/connection/__init__.py chainerpruner/graph/chainer/__init__.py

@@ -5,11 +5,11 @@
 import networkx as nx
 
 import chainer
-from chainerpruner.trace import TraceLinkHook, TraceFunctionHook
-from chainerpruner.node import Node
+from chainerpruner.graph.chainer.trace import TraceLinkHook, TraceFunctionHook
+from chainerpruner.graph.node import Node
 
 
-class Graph():
+class ChainerGraph():
     """Computation Graph Parser
 
     Chainerの計算グラフはLinkとFunctionから構成される

chainerpruner/graph.py chainerpruner/graph/chainer/graph.py

@@ -5,7 +5,7 @@
 import weakref
 import chainer
 
-from chainerpruner.node import Node
+from chainerpruner.graph.node import Node
 
 
 class TraceFunctionHook(chainer.FunctionHook):

chainerpruner/trace.py chainerpruner/graph/chainer/trace.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2018 DeNA Co., Ltd.
+# Licensed under The MIT License [see LICENSE for details]
+
+from chainerpruner import utils
+
+
+class Graph():
+    """Computation Graph Parser
+
+    Chainerの計算グラフはLinkとFunctionから構成される
+    Channel Pruningする上では、現在の層をPruningしたときに
+    後続する層のどこまでその影響があるかを知ることが重要になる
+    ここではモデルをパースして、linksに各層とその影響する層の情報をList[Node]として格納している
+    """
+
+    def __init__(self, model, args):
+        if utils.is_chainer_model(model):
+            from chainerpruner.graph.chainer.graph import ChainerGraph
+            graph = ChainerGraph(model, args)
+            self.is_chainer = True
+            self.is_pytorch = False
+        else:
+            from chainerpruner.graph.pytorch.graph import PyTorchGraph
+            graph = PyTorchGraph(model, args)
+            self.is_chainer = False
+            self.is_pytorch = True
+
+        self.graph = graph.graph
+
+    def plot(self, options=None):
+        import networkx as nx
+        if not options:
+            options = dict()
+        nx.draw(self.graph, with_label=True, **options)

chainerpruner/graph/graph.py

@@ -1,18 +1,14 @@
 # Copyright (c) 2018 DeNA Co., Ltd.
 # Licensed under The MIT License [see LICENSE for details]
 
-from collections import Sequence
-
-import chainer
-
 
 class Node():
     """計算グラフのノード。channel pruningに必要な情報を抽出している。
 
     入出力のテンソルのサイズや、ノードの型、chainerのLinkやFuntionのオブジェクト自体もアクセスできる
     """
 
-    def __init__(self, id_, type_, args, out, link=None, function=None):
+    def __init__(self, id_, type_, args=None, out=None, link=None, function=None, module=None):
         """
 
         nextは、List[Node or List[Node]]
@@ -26,10 +22,10 @@ def __init__(self, id_, type_, args, out, link=None, function=None):
             function:
         """
 
-        if not all([isinstance(a, chainer.variable.VariableNode) for a in args]):
-            raise TypeError()
-        if not all([isinstance(o, chainer.variable.VariableNode) for o in out]):
-            raise TypeError()
+        if not args:
+            args = []
+        if not out:
+            out = []
 
         self.name = id_
         self.id = id_
@@ -40,6 +36,7 @@ def __init__(self, id_, type_, args, out, link=None, function=None):
         self.output_shape = [o.shape if hasattr(o, 'shape') else None for o in out]
         self.link = link
         self.function = function
+        # self.module = module
 
     def __repr__(self):
         if self.link:

chainerpruner/node.py chainerpruner/graph/node.py

@@ -0,0 +1,40 @@
+# Copyright (c) 2018 DeNA Co., Ltd.
+# Licensed under The MIT License [see LICENSE for details]
+
+from typing import Sequence
+import networkx as nx
+
+from chainerpruner.graph.pytorch.trace import ModuleTraceHook
+from chainerpruner.graph.node import Node
+
+
+class PyTorchGraph():
+    """Computation Graph Parser
+    """
+
+    def __init__(self, model, args):
+        # convert to tuple of Variable
+        if not isinstance(args, Sequence):
+            args = [args]
+
+        # 計算グラフを構築しながら、hookを利用して各Moduleの情報を取得する
+        hook = ModuleTraceHook(model)
+        model(*args)
+
+        nodes = hook.graph  # type: Sequence[Node]
+        self.graph = self._traverse_connections(nodes)
+
+    def _traverse_connections(self, nodes):
+        # Hookで得たlinkとfunctionのList[Node]をマージしてグラフにする
+
+        # node.nameをキーにしたGraph
+        # node attrでnodeの実体へアクセス
+        graph = nx.DiGraph()
+
+        for node in nodes:
+            graph.add_node(node)
+            for next_node in nodes:
+                if set(node.output_id) & set(next_node.input_id):
+                    graph.add_edge(node, next_node)
+
+        return graph

chainerpruner/rebuild/links/normalization/__init__.py chainerpruner/graph/pytorch/__init__.py

@@ -0,0 +1,67 @@
+# Copyright (c) 2018 DeNA Co., Ltd.
+# Licensed under The MIT License [see LICENSE for details]
+
+from typing import Sequence
+
+from chainerpruner.graph.node import Node
+
+
+class ModuleTraceHook():
+    """modelをトレースしてpruningに必要な計算グラフの情報を取得する
+
+    トレース後のグラフは`graph`属性として得られる
+
+    """
+
+    def __init__(self, model):
+        self.graph = []
+        # グラフ内各moduleをuniqueに特定するための名前
+        self._module_global_names = dict()
+        self._regist(model)
+
+    def _regist(self, model):
+        for name, module in model.named_modules():
+            # nn.Sequentialやユーザー定義の複数のmoduleをもつmoduleではなく、
+            # Conv2dなどの単一のmoduleのみ解析したいため
+            n_modules = len(list(module.modules()))
+            if n_modules == 1:
+                # print('regist {}'.format(name))
+                self._module_global_names[module] = name
+                module.register_forward_hook(self._hook)
+
+    def _hook(self, module, input, output):
+        """
+
+        Args:
+            module:
+            input: tuple(Tensor)
+            output: Tensor
+
+        Returns:
+
+        """
+        if not isinstance(input, Sequence):
+            input = [input]
+        if not isinstance(output, Sequence):
+            output = [output]
+
+        name = self._module_global_names[module]
+
+        node = Node(id_=name,
+                    type_=type(module),
+                    args=None,
+                    out=None,
+                    # module=module,
+                    link=module,
+                    )
+
+        # TODO(tkat0) mv
+        node.input_id = [i._cdata for i in input]
+        node.input_shape = [tuple(i.shape) if hasattr(i, 'shape') else None for i in input]
+        node.output_id = [o._cdata for o in output if o is not None]
+        node.output_shape = [tuple(o.shape) if hasattr(o, 'shape') else None for o in output]
+
+        if set(node.input_id) == set(node.output_id):
+            pass  # skip in-place op (like ReLU)
+        else:
+            self.graph.append(node)

chainerpruner/graph/pytorch/graph.py

@@ -4,9 +4,8 @@
 import logging
 import numpy as np
 
-from chainer import links as L
-
 from chainerpruner.rebuild import calc_pruning_connection
+from chainerpruner.utils import named_modules
 
 logger = logging.getLogger(__name__)
 
@@ -17,10 +16,11 @@ class Mask():
 
     def __init__(self, model, graph, target_layers, mask_layer=None):
         self.model = model
-        self._model_dict = {name: link for name, link in model.namedlinks()}
+        self._model_dict = {name: link for name, link in named_modules(model)}
         self.graph = graph
         self.target_layers = target_layers
         self.logger = logger
+        self._is_chainer = graph.is_chainer
         self.pruning_connection_info = calc_pruning_connection(graph)
         self.masks = dict()
         self._mask_layer = mask_layer
@@ -31,6 +31,9 @@ def __init__(self, model, graph, target_layers, mask_layer=None):
         if mask_layer not in cand_mask_layer:
             raise AttributeError('mask_layer is expected which {}'.format(cand_mask_layer))
 
+    def is_chainer(self):
+        return self._is_chainer
+
     def get_filter_norm(self, mask):
         """get mask for pruning
 
@@ -58,11 +61,44 @@ def get_thresholds(self, name, mask):
         raise NotImplementedError()
 
     def _get_mask(self, name):
+        """
+
+        Args:
+            name:
+
+        Returns:
+            (NDArray, ndarray): (conv-weight, mask-tensor)
+                conv-weight: (oc, ic, k, k) kernel order
+                mask-tensor: (oc, ic, k, k) or (oc, 1, 1, 1)
+
+        """
+        if self.is_chainer():
+            return self._get_mask_chainer(name)
+        else:
+            return self._get_mask_pytorch(name)
+
+    def _get_mask_pytorch(self, name):
+        from torch import nn
+        conv = self._model_dict[name]
+        if self._mask_layer is None:
+            mask = conv.weight.data.clone()
+        elif self._mask_layer == 'batchnorm':
+            # propagate mask bn: conv-bn
+            post_conv_bn_name = self.pruning_connection_info[name][0]
+            bn = self._model_dict[post_conv_bn_name]
+            if not isinstance(bn, nn.BatchNorm2d):
+                raise ValueError('expected {}(Conv) -> {}(BatchNorm)'.format(name, post_conv_bn_name))
+            mask = bn.weight.data.clone()
+            mask = mask.reshape(-1, 1, 1, 1) # to mask conv weight (oc, ic, kh, kw)
+        return conv.weight.data, mask
+
+    def _get_mask_chainer(self, name):
+        from chainer import links as L
         conv = self._model_dict[name]
         if self._mask_layer is None:
             mask = conv.W.array.copy()
         elif self._mask_layer == 'batchnorm':
-            # conv-bn
+            # propagate mask bn: conv-bn
             post_conv_bn_name = self.pruning_connection_info[name][0]
             bn = self._model_dict[post_conv_bn_name]
             if not isinstance(bn, L.BatchNormalization):
@@ -80,7 +116,8 @@ def __call__(self):
 
         # get mask vector
         target_weights = []
-        for name, link in self.model.namedlinks(skipself=True):
+        options = {'skipself': True} if self.is_chainer() else dict()
+        for name, link in named_modules(self.model, **options):
 
             self.logger.debug('name: %s', name)
 
@@ -93,6 +130,7 @@ def __call__(self):
             out_channels = mask.shape[0]
             mask = self.get_filter_norm(mask)
             if mask.shape != (out_channels, 1, 1, 1):
+                # expected (oc, ic, k, k) kernel order
                 raise RuntimeError()
 
             self.masks[name] = mask
@@ -109,7 +147,10 @@ def __call__(self):
 
             # apply mask
             mask = mask_ >= threshold  # 0: pruning, 1: non-pruning
-            mask = mask.astype(np.float32)
+            try:
+                mask = mask.astype(np.float32)
+            except AttributeError:
+                mask = mask.type_as(target_weight)
 
             info_ = {
                 'name': name,