Pytorch Parser

mmdnn/conversion/_script/convertToIR.py

@@ -71,8 +71,9 @@ def _convert(args):
 
     elif args.srcFramework == 'pytorch':
         assert args.inputShape != None
-        from mmdnn.conversion.pytorch.pytorch_parser import PyTorchParser
-        parser = PyTorchParser(args.network, args.inputShape)
+        inputshape = [int(x) for x in args.inputShape]
+        from mmdnn.conversion.pytorch.pytorch_parser import PytorchParser
+        parser = PytorchParser(args.network, inputshape)
 
     elif args.srcFramework == 'torch' or args.srcFramework == 'torch7':
         from mmdnn.conversion.torch.torch_parser import TorchParser
@@ -108,7 +109,7 @@ def _get_parser():
     parser.add_argument(
         '--srcFramework', '-f',
         type=_text_type,
-        choices=["caffe", "caffe2", "cntk", "mxnet", "keras", "tensorflow", 'tf', 'torch', 'torch7', 'onnx', 'darknet', 'coreml'],
+        choices=["caffe", "caffe2", "cntk", "mxnet", "keras", "tensorflow", 'tf', 'torch', 'torch7', 'onnx', 'darknet', 'coreml', 'pytorch'],
         help="Source toolkit name of the model to be converted.")
 
     parser.add_argument(

mmdnn/conversion/caffe/caffe_emitter.py

@@ -154,7 +154,7 @@ def _shapeToStr(shapes):
 
     def check_if_need_transpose(self, IR_node):
         parent = self.IR_graph.get_parent(IR_node.name, [0])
-        while parent.type == 'Flatten':
+        while parent.type == 'Flatten' or parent.type == 'Dropout':
             parent = self.IR_graph.get_parent(parent.name, [0])
         dim = len(parent.layer.attr['_output_shapes'].list.shape[0].dim)
         if dim > 2:
@@ -213,6 +213,7 @@ def emit_Conv(self, IR_node):
 
     def compute_output_shape(self, IR_node, kernel_h, kernel_w):
         parent_node = self.IR_graph.get_parent(IR_node.name, [0])
+
         if parent_node.get_attr('_output_shapes'):
             shape = parent_node.get_attr('_output_shapes')[0]
             shape = shape_to_list(shape)

mmdnn/conversion/examples/imagenet_test.py

@@ -138,26 +138,26 @@ class TestKit(object):
         },
 
         'pytorch' : {
-            'alexnet'       : lambda path : TestKit.Normalize(path, 227),
-            'densenet121'   : lambda path : TestKit.Normalize(path, 224),
-            'densenet169'   : lambda path : TestKit.Normalize(path, 224),
-            'densenet161'   : lambda path : TestKit.Normalize(path, 224),
-            'densenet201'   : lambda path : TestKit.Normalize(path, 224),
-            'vgg11'         : lambda path : TestKit.Normalize(path, 224),
-            'vgg13'         : lambda path : TestKit.Normalize(path, 224),
-            'vgg16'         : lambda path : TestKit.Normalize(path, 224),
-            'vgg19'         : lambda path : TestKit.Normalize(path, 224),
-            'vgg11_bn'         : lambda path : TestKit.Normalize(path, 224),
-            'vgg13_bn'         : lambda path : TestKit.Normalize(path, 224),
-            'vgg16_bn'         : lambda path : TestKit.Normalize(path, 224),
-            'vgg19_bn'         : lambda path : TestKit.Normalize(path, 224),
-            'resnet18'      : lambda path : TestKit.Normalize(path, 224),
-            'resnet34'      : lambda path : TestKit.Normalize(path, 224),
-            'resnet50'      : lambda path : TestKit.Normalize(path, 224),
-            'resnet101'      : lambda path : TestKit.Normalize(path, 224),
-            'resnet152'     : lambda path : TestKit.Normalize(path, 224),
-            'squeezenet1_0' : lambda path : TestKit.Normalize(path, 224),
-            'inception_v3'  : lambda path : TestKit.Normalize(path, 299),
+            'alexnet'       : lambda path : TestKit.Standard(path, 227),
+            'densenet121'   : lambda path : TestKit.Standard(path, 224),
+            'densenet169'   : lambda path : TestKit.Standard(path, 224),
+            'densenet161'   : lambda path : TestKit.Standard(path, 224),
+            'densenet201'   : lambda path : TestKit.Standard(path, 224),
+            'vgg11'         : lambda path : TestKit.Standard(path, 224),
+            'vgg13'         : lambda path : TestKit.Standard(path, 224),
+            'vgg16'         : lambda path : TestKit.Standard(path, 224),
+            'vgg19'         : lambda path : TestKit.Standard(path, 224),
+            'vgg11_bn'         : lambda path : TestKit.Standard(path, 224),
+            'vgg13_bn'         : lambda path : TestKit.Standard(path, 224),
+            'vgg16_bn'         : lambda path : TestKit.Standard(path, 224),
+            'vgg19_bn'         : lambda path : TestKit.Standard(path, 224),
+            'resnet18'      : lambda path : TestKit.Standard(path, 224),
+            'resnet34'      : lambda path : TestKit.Standard(path, 224),
+            'resnet50'      : lambda path : TestKit.Standard(path, 224),
+            'resnet101'      : lambda path : TestKit.Standard(path, 224),
+            'resnet152'     : lambda path : TestKit.Standard(path, 224),
+            'squeezenet1_0' : lambda path : TestKit.Standard(path, 224),
+            'inception_v3'  : lambda path : TestKit.Standard(path, 299),
         },
 
         'cntk' : {

mmdnn/conversion/mxnet/mxnet_emitter.py

@@ -322,7 +322,7 @@ def emit_FullyConnected(self, IR_node):
         if self.weight_loaded:
             weight_dict = self.weights[IR_node.name]
             parent = self.IR_graph.get_parent(IR_node.name, [0])
-            while parent.type == "Flatten":
+            while parent.type == "Flatten" or parent.type == 'Dropout':
                 parent = self.IR_graph.get_parent(parent.name, [0])
             dim = len(parent.layer.attr['_output_shapes'].list.shape[0].dim)
             if dim > 2:

mmdnn/conversion/pytorch/README.md

@@ -1,26 +1,42 @@
 # PyTorch README
 
-Currently, we only implemented the IR -> PyTorch part.
+Currently, we have already implemented both the the PyTorch -> IR part and the IR -> PyTorch part.
 
-I am implementing the PyTorch parser in branch [pytorch](https://github.com/Microsoft/MMdnn/tree/pytorch) and have some issue with getting layer shape and jit CppOP. Waiting for JIT completion.
+The PyTorch parser is modified from branch [pytorch](https://github.com/Microsoft/MMdnn/tree/pytorch) , using jit CppOP to build the graph.
 
-Any contribution to PyTorch parser (PyTorch -> IR) is welcome.
+Any contribution is welcome.
 
 ## Extract PyTorch pre-trained models
 
 You can refer [PyTorch model extractor](https://github.com/Microsoft/MMdnn/blob/master/mmdnn/conversion/examples/pytorch/extractor.py) to extract your pytorch models.
 
 ```bash
 $ mmdownload -f pytorch -h
-Support frameworks: ['alexnet', 'densenet121', 'densenet161', 'densenet169', 'densenet201', 'inception_v3', 'resnet101', 'resnet152', 'resnet18', 'resnet34', 'resnet50', 'squeezenet1_0', 'squeezenet1_1', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn', 'vgg19', 'vgg19_bn']
+Support frameworks: ['alexnet', 'densenet121', 'densenet161', 'densenet169', 'densenet201', 'inception_v3', 'resnet101', 'resnet152', 'resnet18', 'resnet34', 'resnet50', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn', 'vgg19', 'vgg19_bn']
 
 $ mmdownload -f pytorch -n resnet50 -o ./
 Downloading: "https://download.pytorch.org/models/resnet50-19c8e357.pth" to /home/ruzhang/.torch/models/resnet50-19c8e357.pth
 100%|████████████████████████████████████████████████████████████████████████| 102502400/102502400 [00:06<00:00, 15858546.50it/s]
-PyTorch pretrained model is saved as [.//imagenet_resnet50.pth].
+PyTorch pretrained model is saved as [./imagenet_resnet50.pth].
 
 ```
 
+### Convert Pytorch pre-trained models to IR
+You can convert the whole pytorch model to IR structure. Please remember for the generality, we now only take the whole model `pth`, not just the state dict. To be more specific, it is save using `torch.save()` and `torch.load()` can load the whole model.
+
+```bash
+$ mmtoir -f pytorch -d resnet50 --inputShape 3 224 224 -n imagenet_resnet50.pth --dstNodeName MMdnn_Output
+```
+
+Please bear in mind that always add `--inputShape` argparse. This thing is different from other framework because pytorch is a dynamic framework.
+
+Then you will get
+```
+IR network structure is saved as [resnet50.json].
+IR network structure is saved as [resnet50.pb].
+IR weights are saved as [resnet50.npy].
+```
+
 ### Convert models from IR to PyTorch code snippet and weights
 
 You can use following bash command to convert the IR architecture file [*inception_v3.pb*] and weights file [*inception_v3.npy*] to Caffe Python code file[*pytorch_inception_v3.py*] and IR weights file suit for caffe model[*pytorch_inception_v3.npy*]
@@ -57,6 +73,10 @@ Ubuntu 16.04 with
 
 @ 2018/04/25
 
+## Links
+
+- [pytorch to keras converter](https://github.com/nerox8664/pytorch2keras)
+
 ## Limitation
 
 - The main dataflow in a pytorch network is converted from NHWC(channel last) to NCHW(channel first) format, but some operators (like Concat) with axis may not transform correctly. You may need to correct it manually.

mmdnn/conversion/pytorch/pytorch_parser.py

@@ -166,6 +166,34 @@ def gen_Input(self):
             new_dim = IR_node.attr["shape"].shape.dim.add()
             new_dim.size = self.input_shape[index]
 
+        shape = graph_pb2.TensorShape()
+        new_dim = shape.dim.add()
+        shape_pytorch = self.input_shape
+
+        if len(shape_pytorch) == 4:
+
+            if shape_pytorch[0] == 1:
+                new_dim.size = -1
+            else:
+                new_dim.size = shape_pytorch[0]
+            for index in [2, 3, 1]:
+                new_dim = shape.dim.add()
+                dim = shape_pytorch[index]
+                new_dim.size = dim if dim else -1
+        elif len(shape_pytorch) == 2:
+            if shape_pytorch[0] == 1:
+                new_dim.size = -1
+            else:
+                new_dim.size = shape_pytorch[0]
+            for _ in range(2):
+                new_dim = shape.dim.add()
+                new_dim.size = 1
+            new_dim = shape.dim.add()
+            dim = shape_pytorch[1]
+            new_dim.size = dim if dim else -1
+
+
+        IR_node.attr["_output_shapes"].list.shape.extend([shape])
 
 
     def rename_Conv(self, source_node):

mmdnn/conversion/tensorflow/tensorflow_emitter.py

@@ -177,20 +177,48 @@ def emit_Pool(self, IR_node):
                 IR_node.name))
 
         else:
-            kernel_shape_str = ', '.join('%s' % i for i in IR_node.get_attr('kernel_shape'))
-            strides_str = ', '.join('%s' % i for i in IR_node.get_attr('strides'))
+            dim = len(IR_node.get_attr("strides")) - 2
+            dilations = IR_node.get_attr('dilations')
+            if dilations:
+                for e in IR_node.get_attr('dilations'):
+                    assert e == 1
 
-            input_node, padding = self._defuse_padding(IR_node, padding_const)
+            pool_size = IR_node.get_attr('kernel_shape')[1:-1]
 
-            self.add_body(1, "{:<15} = tf.nn.{}{}({}, [{}], [{}], padding='{}', name='{}')".format(
-                IR_node.variable_name,
-                op,
-                dim_str,
-                input_node,
-                kernel_shape_str,
-                strides_str,
-                padding,
-                IR_node.name))
+            strides = IR_node.get_attr('strides')[1:-1]
+            padding = IR_node.get_attr('pads')[1:dim]
+
+            if pooling_type == "AVG" and pool_size.count(pool_size[0]) == len(pool_size) and strides[0] == 1 and strides.count(strides[0]) == len(strides) and padding.count(padding[0]) == len(padding) and pool_size[0] == padding[0]*2 + 1:
+                kernel_shape_str = ', '.join('%s' % i for i in IR_node.get_attr('kernel_shape'))
+                strides_str = ', '.join('%s' % i for i in IR_node.get_attr('strides'))
+
+
+                self.add_body(1, "{:<15} = tf.nn.{}{}({}, [{}], [{}], padding='{}', name='{}')".format(
+                    IR_node.variable_name,
+                    op,
+                    dim_str,
+                    self.parent_variable_name(IR_node),
+                    kernel_shape_str,
+                    strides_str,
+                    'SAME',
+                    IR_node.name))
+
+            else:
+
+                kernel_shape_str = ', '.join('%s' % i for i in IR_node.get_attr('kernel_shape'))
+                strides_str = ', '.join('%s' % i for i in IR_node.get_attr('strides'))
+
+                input_node, padding = self._defuse_padding(IR_node, padding_const)
+
+                self.add_body(1, "{:<15} = tf.nn.{}{}({}, [{}], [{}], padding='{}', name='{}')".format(
+                    IR_node.variable_name,
+                    op,
+                    dim_str,
+                    input_node,
+                    kernel_shape_str,
+                    strides_str,
+                    padding,
+                    IR_node.name))
 
 
     def emit_UNKNOWN(self, IR_node):

tests/test_conversion_imagenet.py

@@ -789,25 +789,25 @@ def OnnxEmit(original_framework, architecture_name, architecture_path, weight_pa
         },
 
         'pytorch' : {
-            'alexnet'     : [KerasEmit, PytorchEmit],
-            'densenet121' : [KerasEmit, PytorchEmit],
-            'densenet169' : [KerasEmit, PytorchEmit],
-            'densenet161' : [KerasEmit, PytorchEmit],
-            'densenet201' : [KerasEmit, PytorchEmit],
-            'inception_v3': [KerasEmit, PytorchEmit],
-            'vgg11'       : [KerasEmit, PytorchEmit],
-            'vgg13'       : [KerasEmit, PytorchEmit],
-            'vgg16'       : [KerasEmit, PytorchEmit],
-            'vgg19'       : [KerasEmit, PytorchEmit],
-            'vgg11_bn'    : [KerasEmit, PytorchEmit],
-            'vgg13_bn'    : [KerasEmit, PytorchEmit],
-            'vgg16_bn'    : [KerasEmit, PytorchEmit],
-            'vgg19_bn'    : [KerasEmit, PytorchEmit],
-            'resnet18'    : [KerasEmit, PytorchEmit],
-            'resnet34'    : [KerasEmit, PytorchEmit],
-            'resnet50'    : [KerasEmit, PytorchEmit],
-            'resnet101'   : [KerasEmit, PytorchEmit],
-            'resnet152'   : [KerasEmit, PytorchEmit],
+            'alexnet'     : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
+            'densenet121' : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
+            'densenet169' : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
+            'densenet161' : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
+            'densenet201' : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
+            'inception_v3': [CaffeEmit, CoreMLEmit, KerasEmit, PytorchEmit, TensorflowEmit],  # Mxnet broken https://github.com/apache/incubator-mxnet/issues/10194
+            'vgg11'       : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
+            'vgg13'       : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
+            'vgg16'       : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
+            'vgg19'       : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
+            'vgg11_bn'    : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
+            'vgg13_bn'    : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
+            'vgg16_bn'    : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
+            'vgg19_bn'    : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
+            'resnet18'    : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
+            'resnet34'    : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
+            'resnet50'    : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
+            'resnet101'   : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
+            'resnet152'   : [CaffeEmit, CoreMLEmit, KerasEmit, MXNetEmit, PytorchEmit, TensorflowEmit],
         }
 
     }