cntk parser and MXNetEmit

mmdnn/conversion/examples/imagenet_test.py

@@ -64,7 +64,7 @@ class TestKit(object):
         'caffe' : {
             'alexnet'       : lambda path : TestKit.ZeroCenter(path, 227, True),
             'vgg19'         : lambda path : TestKit.ZeroCenter(path, 224, True),
-            'inception_v1'  : lambda path : TestKit.ZeroCenter(path, 227, True),
+            'inception_v1'  : lambda path : TestKit.ZeroCenter(path, 224, True),
             'resnet152'     : lambda path : TestKit.ZeroCenter(path, 224, True),
             'squeezenet'    : lambda path : TestKit.ZeroCenter(path, 227, False)
         },

tests/test_conversion_imagenet.py

@@ -1,4 +1,5 @@
 import os
+import sys
 import six
 import unittest
 import numpy as np
@@ -6,19 +7,11 @@ from six.moves import reload_module
 import tensorflow as tf
 from mmdnn.conversion.examples.imagenet_test import TestKit
 
-from mmdnn.conversion.examples.keras.extractor import keras_extractor
-from mmdnn.conversion.examples.mxnet.extractor import mxnet_extractor
-from mmdnn.conversion.examples.caffe.extractor import caffe_extractor
-from mmdnn.conversion.examples.cntk.extractor import cntk_extractor
-
-from mmdnn.conversion.keras.keras2_parser import Keras2Parser
-from mmdnn.conversion.mxnet.mxnet_parser import MXNetParser
-from mmdnn.conversion.cntk.cntk_parser import CntkParser
-
 from mmdnn.conversion.cntk.cntk_emitter import CntkEmitter
 from mmdnn.conversion.tensorflow.tensorflow_emitter import TensorflowEmitter
 from mmdnn.conversion.keras.keras2_emitter import Keras2Emitter
 from mmdnn.conversion.pytorch.pytorch_emitter import PytorchEmitter
+from mmdnn.conversion.mxnet.mxnet_emitter import MXNetEmitter
 
 def _compute_SNR(x,y):
     noise = x - y
@@ -66,7 +59,7 @@ class CorrectnessTest(unittest.TestCase):
         print("PSNR:", PSNR)
         # self.assertGreater(SNR, self.snr_thresh)
         # self.assertGreater(PSNR, self.psnr_thresh)
-        # self.assertLess(error, self.err_thresh)
+        self.assertLess(error, self.err_thresh)
 
 
 class TestModels(CorrectnessTest):
@@ -82,41 +75,53 @@ class TestModels(CorrectnessTest):
 
         # get original model prediction result
         original_predict = tensorflow_extractor.inference(architecture_name, TestModels.cachedir, image_path)
+        del tensorflow_extractor
 
         # original to IR
+        IR_file = TestModels.tmpdir + 'tensorflow_' + architecture_name + "_converted"
         parser = TensorflowParser(
             TestModels.cachedir + "imagenet_" + architecture_name + ".ckpt.meta",
             TestModels.cachedir + "imagenet_" + architecture_name + ".ckpt",
             None,
             "MMdnn_Output")
-        parser.run(TestModels.tmpdir + architecture_name + "_converted")
+        parser.run(IR_file)
         del parser
         del TensorflowParser
-        del tensorflow_extractor
+
         return original_predict
 
 
     @staticmethod
     def KerasParse(architecture_name, image_path):
+        from mmdnn.conversion.examples.keras.extractor import keras_extractor
+        from mmdnn.conversion.keras.keras2_parser import Keras2Parser
+
         # get original model prediction result
         original_predict = keras_extractor.inference(architecture_name, TestModels.cachedir, image_path)
 
         # download model
         model_filename = keras_extractor.download(architecture_name, TestModels.cachedir)
+        del keras_extractor
 
         # original to IR
+        IR_file = TestModels.tmpdir + 'keras_' + architecture_name + "_converted"
         parser = Keras2Parser(model_filename)
-        parser.run(TestModels.tmpdir + architecture_name + "_converted")
+        parser.run(IR_file)
         del parser
+        del Keras2Parser
         return original_predict
 
     @staticmethod
     def MXNetParse(architecture_name, image_path):
+        from mmdnn.conversion.examples.mxnet.extractor import mxnet_extractor
+        from mmdnn.conversion.mxnet.mxnet_parser import MXNetParser
+
         # download model
         architecture_file, weight_file = mxnet_extractor.download(architecture_name, TestModels.cachedir)
 
         # get original model prediction result
         original_predict = mxnet_extractor.inference(architecture_name, TestModels.cachedir, image_path)
+        del mxnet_extractor
 
         # original to IR
         import re
@@ -125,20 +130,24 @@ class TestModels(CorrectnessTest):
         prefix, epoch = weight_file.rsplit('-', 1)
         model = (architecture_file, prefix, epoch, [3, 224, 224])
 
+        IR_file = TestModels.tmpdir + 'mxnet_' + architecture_name + "_converted"
         parser = MXNetParser(model)
-        parser.run(TestModels.tmpdir + architecture_name + "_converted")
+        parser.run(IR_file)
         del parser
+        del MXNetParser
 
         return original_predict
 
     @staticmethod
     def CaffeParse(architecture_name, image_path):
+        from mmdnn.conversion.examples.caffe.extractor import caffe_extractor
+
         # download model
         architecture_file, weight_file = caffe_extractor.download(architecture_name, TestModels.cachedir)
 
         # get original model prediction result
-
         original_predict = caffe_extractor.inference(architecture_name,architecture_file, weight_file, image_path)
+        del caffe_extractor
 
         # original to IR
         from mmdnn.conversion.caffe.transformer import CaffeTransformer
@@ -149,13 +158,14 @@ class TestModels(CorrectnessTest):
         from mmdnn.conversion.caffe.writer import ModelSaver, PyWriter
 
         prototxt = graph.as_graph_def().SerializeToString()
-        pb_path = TestModels.tmpdir + architecture_name + "_converted.pb"
+        IR_file = TestModels.tmpdir + 'caffe_' + architecture_name + "_converted"
+        pb_path = IR_file + '.pb'
         with open(pb_path, 'wb') as of:
             of.write(prototxt)
         print ("IR network structure is saved as [{}].".format(pb_path))
 
         import numpy as np
-        npy_path = TestModels.tmpdir + architecture_name + "_converted.npy"
+        npy_path = IR_file + '.npy'
         with open(npy_path, 'wb') as of:
             np.save(of, data)
         print ("IR weights are saved as [{}].".format(npy_path))
@@ -164,51 +174,55 @@ class TestModels(CorrectnessTest):
 
     @staticmethod
     def CntkParse(architecture_name, image_path):
+        from mmdnn.conversion.examples.cntk.extractor import cntk_extractor
+        from mmdnn.conversion.cntk.cntk_parser import CntkParser
         # download model
         architecture_file = cntk_extractor.download(architecture_name, TestModels.cachedir)
 
         # get original model prediction result
         original_predict = cntk_extractor.inference(architecture_name, architecture_file, image_path)
+        del cntk_extractor
 
         # original to IR
+        IR_file = TestModels.tmpdir + 'cntk_' + architecture_name + "_converted"
         parser = CntkParser(architecture_file)
-        parser.run(TestModels.tmpdir + architecture_name + "_converted")
+        parser.run(IR_file)
         del parser
+        del CntkParser
         return original_predict
 
     @staticmethod
     def CntkEmit(original_framework, architecture_name, architecture_path, weight_path, image_path):
         # IR to code
+        converted_file = original_framework + '_cntk_' + architecture_name + "_converted"
+        converted_file = converted_file.replace('.', '_')
         emitter = CntkEmitter((architecture_path, weight_path))
-        emitter.run("converted_model.py", None, 'test')
+        emitter.run(converted_file + '.py', None, 'test')
         del emitter
 
-        # import converted model
+        model_converted = __import__(converted_file).KitModel(weight_path)
-        import converted_model
-        reload_module (converted_model)
-        model_converted = converted_model.KitModel(TestModels.tmpdir + architecture_name + "_converted.npy")
 
         func = TestKit.preprocess_func[original_framework][architecture_name]
         img = func(image_path)
         predict = model_converted.eval({model_converted.arguments[0]:[img]})
         converted_predict = np.squeeze(predict)
         del model_converted
-        del converted_model
+        del sys.modules[converted_file]
-        os.remove("converted_model.py")
+        os.remove(converted_file + '.py')
         return converted_predict
 
 
     @staticmethod
     def TensorflowEmit(original_framework, architecture_name, architecture_path, weight_path, image_path):
         # IR to code
+        converted_file = original_framework + '_tensorflow_' + architecture_name + "_converted"
+        converted_file = converted_file.replace('.', '_')
         emitter = TensorflowEmitter((architecture_path, weight_path))
-        emitter.run("converted_model.py", None, 'test')
+        emitter.run(converted_file + '.py', None, 'test')
         del emitter
 
         # import converted model
-        import converted_model
+        model_converted = __import__(converted_file).KitModel(weight_path)
-        reload_module (converted_model)
-        model_converted = converted_model.KitModel(TestModels.tmpdir + architecture_name + "_converted.npy")
         input_tf, model_tf = model_converted
 
         func = TestKit.preprocess_func[original_framework][architecture_name]
@@ -219,8 +233,8 @@ class TestModels(CorrectnessTest):
             sess.run(init)
             predict = sess.run(model_tf, feed_dict = {input_tf : input_data})
         del model_converted
-        del converted_model
+        del sys.modules[converted_file]
-        os.remove("converted_model.py")
+        os.remove(converted_file + '.py')
         converted_predict = np.squeeze(predict)
         return converted_predict
 
@@ -233,13 +247,11 @@ class TestModels(CorrectnessTest):
         converted_file = original_framework + '_pytorch_' + architecture_name + "_converted"
         converted_file = converted_file.replace('.', '_')
         emitter = PytorchEmitter((architecture_path, weight_path))
-        emitter.run("converted_model.py", "pytorch_weight.npy", 'test')
+        emitter.run(converted_file + '.py', converted_file + '.npy', 'test')
         del emitter
 
         # import converted model
-        import converted_model
+        model_converted = __import__(converted_file).KitModel(converted_file + '.npy')
-        reload_module (converted_model)
-        model_converted = converted_model.KitModel("pytorch_weight.npy")
         model_converted.eval()
 
         func = TestKit.preprocess_func[original_framework][architecture_name]
@@ -253,10 +265,10 @@ class TestModels(CorrectnessTest):
         predict = predict.data.numpy()
 
         del model_converted
-        del converted_model
+        del sys.modules[converted_file]
         del torch
-        os.remove("converted_model.py")
+        os.remove(converted_file + '.py')
-        os.remove("pytorch_weight.npy")
+        os.remove(converted_file + '.npy')
         converted_predict = np.squeeze(predict)
         return converted_predict
 
@@ -264,14 +276,14 @@ class TestModels(CorrectnessTest):
     @staticmethod
     def KerasEmit(original_framework, architecture_name, architecture_path, weight_path, image_path):
         # IR to code
+        converted_file = original_framework + '_keras_' + architecture_name + "_converted"
+        converted_file = converted_file.replace('.', '_')
         emitter = Keras2Emitter((architecture_path, weight_path))
-        emitter.run("converted_model.py", None, 'test')
+        emitter.run(converted_file + '.py', None, 'test')
         del emitter
 
         # import converted model
-        import converted_model
+        model_converted = __import__(converted_file).KitModel(weight_path)
-        reload_module (converted_model)
-        model_converted = converted_model.KitModel(TestModels.tmpdir + architecture_name + "_converted.npy")
 
         func = TestKit.preprocess_func[original_framework][architecture_name]
         img = func(image_path)
@@ -281,59 +293,69 @@ class TestModels(CorrectnessTest):
         converted_predict = np.squeeze(predict)
 
         del model_converted
-        del converted_model
+        del sys.modules[converted_file]
 
         import keras.backend as K
         K.clear_session()
 
-        os.remove("converted_model.py")
+        os.remove(converted_file + '.py')
         return converted_predict
 
     @staticmethod
     def MXNetEmit(original_framework, architecture_name, architecture_path, weight_path, image_path):
+        from collections import namedtuple
+        Batch = namedtuple('Batch', ['data'])
+
+        import mxnet as mx
         print("Testing {} from {} to MXNet.".format(architecture_name, original_framework))
 
-        # # IR to code
+        # IR to code
-        # emitter = Keras2Emitter((architecture_path, weight_path))
+        converted_file = original_framework + '_mxnet_' + architecture_name + "_converted"
-        # emitter.run("converted_model.py", None, 'test')
+        converted_file = converted_file.replace('.', '_')
-        # del emitter
+        output_weights_file = converted_file + "-0000.params"
+        emitter = MXNetEmitter((architecture_path, weight_path, output_weights_file))
+        emitter.run(converted_file + '.py', None, 'test')
+        del emitter
 
-        # # import converted model
+        # import converted model
-        # import converted_model
+        imported = __import__(converted_file)
-        # reload_module (converted_model)
+        model_converted = imported.RefactorModel()
-        # model_converted = converted_model.KitModel(TestModels.tmpdir + architecture_name + "_converted.npy")
+        model_converted = imported.deploy_weight(model_converted, output_weights_file)
 
-        # func = TestKit.preprocess_func[original_framework][architecture_name]
+        func = TestKit.preprocess_func[original_framework][architecture_name]
-        # img = func(image_path)
+        img = func(image_path)
-        # input_data = np.expand_dims(img, 0)
+        img = np.transpose(img, (2, 0, 1))
+        input_data = np.expand_dims(img, 0)
 
-        # predict = model_converted.predict(input_data)
+        model_converted.forward(Batch([mx.nd.array(input_data)]))
-        # converted_predict = np.squeeze(predict)
+        predict = model_converted.get_outputs()[0].asnumpy()
+        converted_predict = np.squeeze(predict)
 
-        # del model_converted
+        del model_converted
-        # del converted_model
+        del sys.modules[converted_file]
+        del mx
 
-        # import keras.backend as K
+        os.remove(converted_file + '.py')
-        # K.clear_session()
+        os.remove(output_weights_file)
-
+        return converted_predict
-        # os.remove("converted_model.py")
-        # return converted_predict
 
     test_table = {
         'cntk' : {
             # 'alexnet'       : [TensorflowEmit, KerasEmit],
             # 'resnet18'      : [TensorflowEmit, KerasEmit],
-            'inception_v3'  : [CntkEmit],
+            'inception_v3'  : [PytorchEmit ],
         },
+
         'keras' : {
-            'vgg16'        : [CntkEmit, TensorflowEmit, KerasEmit, PytorchEmit],
+            # 'vgg16'        : [CntkEmit, TensorflowEmit, KerasEmit, PytorchEmit],
-            'vgg19'        : [CntkEmit, TensorflowEmit, KerasEmit, PytorchEmit],
+            'vgg19'        : [MXNetEmit],
-            'inception_v3' : [CntkEmit, TensorflowEmit, KerasEmit, PytorchEmit],
+            # 'vgg19'        : [CntkEmit, TensorflowEmit, KerasEmit, PytorchEmit],
-            'resnet50'     : [CntkEmit, TensorflowEmit, KerasEmit, PytorchEmit],
+            # 'inception_v3' : [CntkEmit, TensorflowEmit, KerasEmit, PytorchEmit],
-            'densenet'     : [CntkEmit, TensorflowEmit, KerasEmit, PytorchEmit],
+            # 'resnet50'     : [CntkEmit, TensorflowEmit, KerasEmit, PytorchEmit],
-            'xception'     : [TensorflowEmit, KerasEmit],
+            # 'densenet'     : [CntkEmit, TensorflowEmit, KerasEmit, PytorchEmit],
-            'mobilenet'    : [TensorflowEmit, KerasEmit],
+            # 'xception'     : [TensorflowEmit, KerasEmit],
-            'nasnet'       : [TensorflowEmit, KerasEmit],
+            # 'mobilenet'    : [TensorflowEmit, KerasEmit],
+            # 'nasnet'       : [TensorflowEmit, KerasEmit],
         },
 
         'mxnet' : {
@@ -370,49 +392,53 @@ class TestModels(CorrectnessTest):
         ensure_dir(self.tmpdir)
 
         for network_name in self.test_table[original_framework].keys():
-            print("Test {} from {} start.".format(network_name, original_framework), file=sys.stderr, flush=True)
+            # print("Test {} from {} start.".format(network_name, original_framework), file=sys.stderr, flush=True)
+            print("Test {} from {} start.".format(network_name, original_framework))
 
             # get original model prediction result
             original_predict = parser(network_name, self.image_path)
 
+            IR_file = TestModels.tmpdir + original_framework + '_' + network_name + "_converted"
             for emit in self.test_table[original_framework][network_name]:
-                print('Testing conversion {} from {} to {}.'.format(network_name, original_framework, emit.__func__.__name__[:-4]), file=sys.stderr, flush=True)
+                # print('Testing conversion {} from {} to {}.'.format(network_name, original_framework, emit.__func__.__name__[:-4]), file=sys.stderr, flush=True)
+                print('Testing conversion {} from {} to {}.'.format(network_name, original_framework, emit.__func__.__name__[:-4]))
                 converted_predict = emit.__func__(
                     original_framework,
                     network_name,
-                    self.tmpdir + network_name + "_converted.pb",
+                    IR_file + ".pb",
-                    self.tmpdir + network_name + "_converted.npy",
+                    IR_file + ".npy",
                     self.image_path)
 
                 self._compare_outputs(original_predict, converted_predict)
 
-                print('Conversion {} from {} to {} passed.'.format(network_name, original_framework, emit.__func__.__name__[:-4]), file=sys.stderr, flush=True)
+                # print('Conversion {} from {} to {} passed.'.format(network_name, original_framework, emit.__func__.__name__[:-4]), file=sys.stderr, flush=True)
-
+                print('Conversion {} from {} to {} passed.'.format(network_name, original_framework, emit.__func__.__name__[:-4]))
             try:
-                os.remove(self.tmpdir + network_name + "_converted.json")
+                os.remove(IR_file + ".json")
             except OSError:
                 pass
 
-            os.remove(self.tmpdir + network_name + "_converted.pb")
+            os.remove(IR_file + ".pb")
-            os.remove(self.tmpdir + network_name + "_converted.npy")
+            os.remove(IR_file + ".npy")
             print("Testing {} model {} passed.".format(original_framework, network_name))
 
         print("Testing {} model all passed.".format(original_framework))
 
-    def test_cntk(self):
+    # def test_cntk(self):
-        self._test_function('cntk', self.CntkParse)
+    #     self._test_function('cntk', self.CntkParse)
-
-    def test_tensorflow(self):
-        self._test_function('tensorflow', self.TensorFlowParse)
 
 
-    def test_caffe(self):
+    # def test_tensorflow(self):
-        self._test_function('caffe', self.CaffeParse)
+    #     self._test_function('tensorflow', self.TensorFlowParse)
+
+
+    # def test_caffe(self):
+    #     self._test_function('caffe', self.CaffeParse)
 
 
     def test_keras(self):
         self._test_function('keras', self.KerasParse)
 
 
-    def test_mxnet(self):
+    # def test_mxnet(self):
-        self._test_function('mxnet', self.MXNetParse)
+    #     self._test_function('mxnet', self.MXNetParse)