Implement pytorch model extractor.

Author: kitstar

mmdnn/conversion/examples/imagenet_test.py

@@ -6,16 +6,9 @@
 from __future__ import absolute_import
 import argparse
 import numpy as np
-import sys
-import os
 from six import text_type as _text_type
-
-# work for tf 1.4 in windows & linux
 from tensorflow.contrib.keras.api.keras.preprocessing import image
 
-# work for tf 1.3 & 1.4 in linux
-# from tensorflow.contrib.keras.python.keras.preprocessing import image
-
 
 class TestKit(object):
 
@@ -44,6 +37,10 @@ class TestKit(object):
             'resnet'         : [(21, 0.84012794), (144, 0.097428247), (23, 0.039757393), (146, 0.010432643), (99, 0.0023797606)],
             'squeezenet'     : [(21, 0.36026478), (128, 0.084114805), (835, 0.07940048), (144, 0.057378717), (749, 0.053491514)],
             'inception_bn'   : [(21, 0.84332663), (144, 0.041747514), (677, 0.021810319), (973, 0.02054958), (115, 0.008529461)]
+        },
+        'pytorch' :{
+            'resnet152' : [(21, 13.080057), (141, 12.32998), (94, 9.8761454), (146, 9.3761511), (143, 8.9194641)],
+            'vgg19'     : [(821, 8.4734678), (562, 8.3472366), (835, 8.2712851), (749, 7.792901), (807, 6.6604013)],
         }
     }
 
@@ -76,6 +73,12 @@ class TestKit(object):
             'resnet'        : lambda path : TestKit.Identity(path, 224, True),
             'squeezenet'    : lambda path : TestKit.ZeroCenter(path, 224, False),
             'inception_bn'  : lambda path : TestKit.Identity(path, 224, False)
+        },
+
+        'pytorch' : {
+            'vgg19'         : lambda path : TestKit.Normalize(path),
+            'resnet152'     : lambda path : TestKit.Normalize(path),
+            'inception_v3'  : lambda path : TestKit.Normalize(path),
         }
     }
 
@@ -122,6 +125,17 @@ def ZeroCenter(path, size, BGRTranspose=False):
         return x
 
 
+    @staticmethod
+    def Normalize(path, size=224, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
+        img = image.load_img(path, target_size=(size, size))
+        x = image.img_to_array(img)
+        x /= 255.0
+        for i in range(0, 3):
+            x[..., i] -= mean[i]
+            x[..., i] /= std[i]
+        return x
+
+
     @staticmethod
     def Standard(path, size):
         img = image.load_img(path, target_size = (size, size))

mmdnn/conversion/examples/pytorch/extract_model.py

@@ -0,0 +1,71 @@
+#----------------------------------------------------------------------------------------------
+#  Copyright (c) Microsoft Corporation. All rights reserved.
+#  Licensed under the MIT License. See License.txt in the project root for license information.
+#----------------------------------------------------------------------------------------------
+
+import argparse
+import os
+from six import text_type as _text_type
+from mmdnn.conversion.examples.imagenet_test import TestKit
+import torch
+import torchvision.models as models
+
+
+NETWORKS_MAP = {
+    'inception_v3'      : lambda : models.inception_v3(pretrained=True),
+    'vgg16'             : lambda : models.vgg16(pretrained=True),
+    'vgg19'             : lambda : models.vgg19(pretrained=True),
+    'resnet152'         : lambda : models.resnet152(pretrained=True),
+    'densenet'          : lambda : models.densenet201(pretrained=True),
+    'squeezenet'        : lambda : models.squeezenet1_1(pretrained=True)
+}
+
+
+def _main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('-n', '--network',
+                        type=_text_type, help='Model Type', required=True,
+                        choices=NETWORKS_MAP.keys())
+
+    parser.add_argument('-i', '--image', type=_text_type, help='Test Image Path')
+
+    args = parser.parse_args()
+
+    file_name = "imagenet_{}.pt".format(args.network)
+    if not os.path.exists(file_name):
+        model = NETWORKS_MAP.get(args.network)
+        model = model()
+        torch.save(model, file_name)
+        print("PyTorch pretrained model is saved as [{}].".format(file_name))
+    else:
+        print("File [{}] existed!".format(file_name))
+        model = torch.load(file_name)
+
+    if args.image:
+        import numpy as np
+        func = TestKit.preprocess_func['pytorch'][args.network]
+        img = func(args.image)
+        img = np.transpose(img, (2, 0, 1))
+        img = np.expand_dims(img, 0).copy()
+        data = torch.from_numpy(img)
+        data = torch.autograd.Variable(data, requires_grad=False)
+
+        model.eval()
+        predict = model(data).data.numpy()
+        predict = np.squeeze(predict)
+        top_indices = predict.argsort()[-5:][::-1]
+        result = [(i, predict[i]) for i in top_indices]
+        print(result)
+
+        # layer_name = 'block2_pool'
+        # intermediate_layer_model = keras.Model(inputs=model.input,
+        #                                  outputs=model.get_layer(layer_name).output)
+        # intermediate_output = intermediate_layer_model.predict(img)
+        # print (intermediate_output)
+        # print (intermediate_output.shape)
+        # print ("%.30f" % np.sum(intermediate_output))
+
+
+if __name__ == '__main__':
+    _main()

mmdnn/conversion/pytorch/inference.py

@@ -0,0 +1,71 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# Licensed under the MIT license. See LICENSE.md file in the project root
+# for full license information.
+# ==============================================================================
+
+import argparse
+import numpy as np
+import sys
+import os
+import caffe
+
+def inference(args):
+    net = caffe.Net(args.network, args.weights, caffe.TEST)
+
+    from tensorflow.contrib.keras.python.keras.preprocessing import image
+    image_path = 'mmdnn/conversion/examples/data/seagull.jpg'
+
+    if args.preprocess == 'vgg':
+        img = image.load_img(image_path, target_size=(224, 224))
+        x = image.img_to_array(img)
+        # Zero-center by mean pixel
+        x = x[..., ::-1]
+        x[..., 0] -= 103.939
+        x[..., 1] -= 116.779
+        x[..., 2] -= 123.68
+
+        x = np.transpose(x, [2, 0, 1])
+
+    elif args.preprocess == 'resnet' or args.preprocess == 'inception':
+        img = image.load_img(image_path, target_size=(224, 224))
+        x = image.img_to_array(img)
+        x /= 255.0
+        x -= 0.5
+        x *= 2.0
+        x = np.transpose(x, [2, 0, 1])
+
+    else:
+        assert False
+
+    x = np.expand_dims(x, 0)
+    net.blobs['data'].data[...] = x
+    predict = np.squeeze(net.forward()['prob'][0])
+
+    test = 'pool1/norm1'
+    immediate_data = net.blobs[test].data[0]
+    print (immediate_data)
+    print (immediate_data.shape)
+    print ("%.30f" % np.sum(np.array(immediate_data)))
+
+    top_indices = predict.argsort()[-5:][::-1]
+    result = [(i, predict[i]) for i in top_indices]
+    print (result)
+
+
+if __name__=='__main__':
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('-p', '--preprocess',
+                        type=str, choices = ["vgg", "resnet", "inception"], help='Model Preprocess Type', required=False, default='vgg')
+
+    parser.add_argument('-n', '--network',
+                        type=str, required=True)
+
+    parser.add_argument('-w', '--weights',
+                        type=str, required=True)
+
+
+    args = parser.parse_args()
+
+    inference(args)