Support automatic UpSampling op

README.md

@@ -33,7 +33,6 @@
 
 The convertor Is not fully automatically, The convertor not 
 
-+ if you wanna convert  upsampling operator , the convertor will convert Upsampling operator to Deconvolution in Caffe , The Deconvolution channels need to be set (in prototxt_basic.py names_output).
 + If you use Flatten Layer ,You need to manually to connect them becasuse the converted compute graph will be divided into two parts.
 + If convert a detection model. You need to remove the anchor process and put it into post process.
 + Usually,If you find that conversion errors, please set the prefix name of you backbone network in mxnet2caffe.py.

find.py

@@ -0,0 +1,23 @@
+from difflib import SequenceMatcher
+import json
+import collections
+
+
+def find_backbone(json_path):
+    with open(json_path) as json_file:
+        jdata = json.load(json_file)
+
+    matches = []
+    for i_node in range(0, len(jdata['nodes']) - 1):
+        node_i1 = jdata['nodes'][i_node]
+        node_i2 = jdata['nodes'][i_node+1]
+        name1 = (node_i1['name'])
+        name2 = (node_i2['name'])
+
+        match = SequenceMatcher(None, name1, name2).find_longest_match(0, name1.find('_'), 0, name2.find('_'))
+        matches.append(name1[match.a: match.a + match.size])
+
+    counter = collections.Counter(matches)
+    final_match = counter.most_common()[0][0]
+
+    return final_match

json2prototxt.py

@@ -1,42 +1,54 @@
-import sys
-import argparse
 import json
 from prototxt_basic import *
 
-parser = argparse.ArgumentParser(description='Convert MXNet jason to Caffe prototxt')
-parser.add_argument('--mx-json',     type=str, default='R50v2/R50v2-symbol.json')
-parser.add_argument('--cf-prototxt', type=str, default='R50v2/R50v2.prototxt')
-args = parser.parse_args()
-
-with open(args.mx_json) as json_file:    
-  jdata = json.load(json_file)
-  print(jdata)
-
-with open(args.cf_prototxt, "w") as prototxt_file:
-  for i_node in range(0,len(jdata['nodes'])):
-    node_i = jdata['nodes'][i_node]
-    if str(node_i['op']) == 'null' and str(node_i['name']) != 'data':
-      continue
-
-    print('{}, \top:{}, name:{} -> {}'.format(i_node,node_i['op'].ljust(20),
-                                        node_i['name'].ljust(30),
-                                        node_i['name']).ljust(20))
-    info = node_i
-
-    info['top'] = info['name']
-    info['bottom'] = []
-    info['params'] = []
-    for input_idx_i in node_i['inputs']:
-      input_i = jdata['nodes'][input_idx_i[0]]
-      if str(input_i['op']) != 'null' or (str(input_i['name']) == 'data'):
-        info['bottom'].append(str(input_i['name']))
-      if str(input_i['op']) == 'null':
-        info['params'].append(str(input_i['name']))
-        if not str(input_i['name']).startswith(str(node_i['name'])):
-          print('           use shared weight -> %s'% str(input_i['name']))
-          info['share'] = True
+def write_prototxt(json_path, prototx_path, input_shape):
+  with open(json_path) as json_file:    
+    jdata = json.load(json_file)
+    print(jdata)
+
+  with open(prototx_path, "w") as prototxt_file:
+    for i_node in range(0,len(jdata['nodes'])):
+      node_i = jdata['nodes'][i_node]
+      if str(node_i['op']) == 'null' and str(node_i['name']) != 'data':
+        continue
 
-    write_node(prototxt_file, info)
+      print('{}, \top:{}, name:{} -> {}'.format(i_node,node_i['op'].ljust(20),
+                                          node_i['name'].ljust(30),
+                                          node_i['name']).ljust(20))
+      info = node_i
+
+      info['top'] = info['name']
+      info['bottom'] = []
+      info['params'] = []
+      for input_idx_i in node_i['inputs']:
+        input_i = jdata['nodes'][input_idx_i[0]]
+        if str(input_i['op']) != 'null' or (str(input_i['name']) == 'data'):
+          info['bottom'].append(str(input_i['name']))
+        if str(input_i['op']) == 'null':
+          info['params'].append(str(input_i['name']))
+          if not str(input_i['name']).startswith(str(node_i['name'])):
+            print('           use shared weight -> %s'% str(input_i['name']))
+            info['share'] = True
+
+      if str(node_i['op']) == 'UpSampling':
+      found = False
+      while not found:
+        inputs = node_i['inputs']
+        for j in range(len(inputs)):
+          if 'attrs' in jdata['nodes'][inputs[j][0]]:
+            if 'num_filter' in jdata['nodes'][inputs[j][0]]['attrs']:
+              info["num_output"] = int(jdata['nodes'][inputs[j][0]]['attrs']['num_filter'])
+              info["group"] = info["num_output"]
+              found = True
+        node_i = jdata['nodes'][inputs[0][0]]
 
-print("*** JSON to PROTOTXT FINISH ***")
+      if str(node_i['op']) == 'data':
+      for char in ['[', ']', '(', ')']:
+        input_shape = input_shape.replace(char, '')
+      input_shape = [int(item) for item in input_shape.split(',')]
+      info["shape"] = input_shape
+
+      write_node(prototxt_file, info)
+
+  print("*** JSON to PROTOTXT FINISH ***")
 

mxnet2caffe.py

@@ -1,4 +1,4 @@
-import sys, argparse
+import argparse
 import mxnet as mx
 import sys
 import os
@@ -11,18 +11,23 @@
     sys.path.append(os.path.join(curr_path, "/Users/yujinke/me/caffe/python"))
     import caffe
 
-
+from find import *
+from json2prototxt import write_prototxt
 
 import time
-import os
 os.environ["CUDA_VISIBLE_DEVICES"] = '4'
 parser = argparse.ArgumentParser(description='Convert MXNet model to Caffe model')
 parser.add_argument('--mx-model',    type=str, default='model_mxnet/face/facega2')
 parser.add_argument('--mx-epoch',    type=int, default=0)
 parser.add_argument('--cf-prototxt', type=str, default='model_caffe/face/facega2.prototxt')
 parser.add_argument('--cf-model',    type=str, default='model_caffe/face/facega2.caffemodel')
+parser.add_argument('--input_shape', type=str, default='1,3,640,640')
 args = parser.parse_args()
 
+# ------------------------------------------
+# Create prototxt
+write_prototxt(args.mx_model + '-symbol.json', args.cf_prototxt, args.input_shape)
+
 # ------------------------------------------
 # Load
 _, arg_params, aux_params = mx.model.load_checkpoint(args.mx_model, args.mx_epoch)
@@ -42,7 +47,8 @@
 print('----------------------------------\n')
 print('VALID KEYS:')
 
-backbone = "hstage1"
+# backbone = "hstage1"
+backbone = find_backbone(args.mx_model + '-symbol.json')
 
 for i_key,key_i in enumerate(all_keys):
 

prototxt_basic.py

@@ -5,7 +5,7 @@
 attrstr = "attrs"
 #attrstr = "param"
 
-names_output = {"rf_c2_upsampling":256 ,"rf_c3_upsampling":256}
+# names_output = {"rf_c2_upsampling":256 ,"rf_c3_upsampling":256}
 #names_output = {"ssh_m2_red_up":32,"ssh_c3_up":32 }
 
 def data(txt_file, info):
@@ -15,7 +15,10 @@ def data(txt_file, info):
   txt_file.write('  type: "Input"\n')
   txt_file.write('  top: "data"\n')
   txt_file.write('  input_param {\n')
-  txt_file.write('    shape: { dim: 1 dim: 3 dim: 512 dim: 1224 }\n') # TODO
+  txt_file.write('    shape: {{ dim: {} dim: {} dim: {} dim: {} }}\n'.format(info['shape'][0],
+                                                                             info['shape'][1],
+                                                                             info['shape'][2],
+                                                                             info['shape'][3]))
   txt_file.write('  }\n')
   txt_file.write('}\n')
   txt_file.write('\n')
@@ -154,13 +157,13 @@ def Upsampling(txt_file, info):
   print(info[attrstr])
   print(info)
   txt_file.write('  convolution_param {\n')
-  txt_file.write('    num_output: %s\n' % names_output[info["name"]])
+  txt_file.write('    num_output: %s\n' % info["num_output"])
   #txt_file.write('    num_output: %s\n' % info[attrstr]['num_filter'])
   txt_file.write('    kernel_size: %d\n' % (2 * scale - scale % 2))  # TODO
   txt_file.write('    stride: %d\n' % scale)
   txt_file.write('    pad: %d\n' % math.ceil((scale - 1)/2.0))  # TODO
   #txt_file.write('    group: %s\n' % info[attrstr]['num_filter'])
-  txt_file.write('    group: %s\n' % names_output[info["name"]])
+  txt_file.write('    group: %s\n' % info["group"])
 
   txt_file.write('    bias_term: false\n')
   txt_file.write('    weight_filler: {\n')