Adds support for emitting darknet models

microsoft · Dec 10, 2018 · 3dd4cbc · 3dd4cbc
1 parent 72529a9
commit 3dd4cbc
Show file tree

Hide file tree

Showing 6 changed files with 463 additions and 60 deletions.
diff --git a/mmdnn/conversion/_script/IRToCode.py b/mmdnn/conversion/_script/IRToCode.py
@@ -55,6 +55,13 @@ def _convert(args):
             raise NotImplementedError("ONNX emitter needs IR weight file")
         else:
             emitter = OnnxEmitter(args.IRModelPath, args.IRWeightPath)
+    elif args.dstFramework == 'darknet':
+        from mmdnn.conversion.darknet.darknet_emitter import DarknetEmitter
+        if args.IRWeightPath is None:
+            emitter = DarknetEmitter(args.IRModelPath)
+        else:
+            assert args.dstWeightPath
+            emitter = DarknetEmitter((args.IRModelPath, args.IRWeightPath))
     else:
         assert False
 
@@ -79,7 +86,7 @@ def _get_parser():
     parser.add_argument(
         '--dstFramework', '-f',
         type=_text_type,
-        choices=['caffe', 'caffe2', 'cntk', 'mxnet', 'keras', 'tensorflow', 'coreml', 'pytorch', 'onnx'],
+        choices=['caffe', 'caffe2', 'cntk', 'mxnet', 'keras', 'tensorflow', 'coreml', 'pytorch', 'onnx', 'darknet'],
         required=True,
         help='Format of model at srcModelPath (default is to auto-detect).')
 

diff --git a/mmdnn/conversion/_script/convert.py b/mmdnn/conversion/_script/convert.py
@@ -29,7 +29,7 @@ def _get_parser():
     parser.add_argument(
         '--dstFramework', '-df',
         type=_text_type,
-        choices=['caffe', 'caffe2', 'cntk', 'mxnet', 'keras', 'tensorflow', 'coreml', 'pytorch', 'onnx'],
+        choices=['caffe', 'caffe2', 'cntk', 'mxnet', 'keras', 'tensorflow', 'coreml', 'pytorch', 'onnx', 'darknet'],
         required=True,
         help='Format of model at srcModelPath (default is to auto-detect).')
     parser.add_argument(

diff --git a/mmdnn/conversion/_script/dump_code.py b/mmdnn/conversion/_script/dump_code.py
@@ -20,6 +20,8 @@ def dump_code(framework, network_filepath, weight_filepath, dump_filepath, dump_
         from mmdnn.conversion.mxnet.saver import save_model
     elif framework == 'pytorch':
         from mmdnn.conversion.pytorch.saver import save_model
+    elif framework == 'darknet':
+        from mmdnn.conversion.darknet.saver import save_model
     elif framework == 'tensorflow':
         from mmdnn.conversion.tensorflow.saver import save_model
         save_model(MainModel, network_filepath, weight_filepath, dump_filepath, dump_tag)
@@ -40,7 +42,7 @@ def _get_parser():
     parser = argparse.ArgumentParser(description='Dump the model code into target model.')
 
     parser.add_argument(
-        '-f', '--framework', type=_text_type, choices=["caffe", "cntk", "mxnet", "keras", "tensorflow", 'pytorch', 'onnx'],
+        '-f', '--framework', type=_text_type, choices=["caffe", "cntk", "mxnet", "keras", "tensorflow", 'pytorch', 'onnx', 'darknet'],
         required=True,
         help='Format of model at srcModelPath (default is to auto-detect).'
     )