Modify the caffe_converter to use the new api. (#4655)

Compatible with python2 / 3
Specification naming

tools/caffe_converter/caffe_parse/parse_from_protobuf.py

@@ -1,34 +1,38 @@
 from google.protobuf import text_format
 import numpy as np
-import caffe_pb2
+import caffe_parse.caffe_pb2 as caffe_pb2
 
-def parse_caffemodel(filepath):
-    '''
+
+def parse_caffemodel(file_path):
+    """
     parses the trained .caffemodel file
 
     filepath: /path/to/trained-model.caffemodel
 
     returns: layers
-    '''
-    f = open(filepath, 'rb')
+    """
+    f = open(file_path, 'rb')
     contents = f.read()
 
-    netparam = caffe_pb2.NetParameter()
-    netparam.ParseFromString(contents)
+    net_param = caffe_pb2.NetParameter()
+    net_param.ParseFromString(contents)
 
-    layers = find_layers(netparam)
+    layers = find_layers(net_param)
     return layers
 
-def find_layers(netparam):
-    if len(netparam.layers) > 0:
-        return netparam.layers
-    elif len(netparam.layer) > 0:
-        return netparam.layer
+
+def find_layers(net_param):
+    if len(net_param.layers) > 0:
+        return net_param.layers
+    elif len(net_param.layer) > 0:
+        return net_param.layer
     else:
-        raise Exception ("Couldn't find layers")
+        raise Exception("Couldn't find layers")
+
 
 def main():
     param_dict = parse_caffemodel('xxx.caffemodel')
 
+
 if __name__ == '__main__':
     main()

tools/caffe_converter/convert_model.py

@@ -1,3 +1,4 @@
+from __future__ import print_function
 import mxnet as mx
 import numpy as np
 import argparse
@@ -9,15 +10,18 @@
     import caffe
 except ImportError:
     import caffe_parse.parse_from_protobuf as parse
+
     caffe_flag = False
 
+
 def get_caffe_iter(layer_names, layers):
     for layer_idx, layer in enumerate(layers):
         layer_name = re.sub('[-/]', '_', layer_names[layer_idx])
         layer_type = layer.type
         layer_blobs = layer.blobs
         yield (layer_name, layer_type, layer_blobs)
 
+
 def get_iter(layers):
     for layer in layers:
         layer_name = re.sub('[-/]', '_', layer.name)
@@ -46,56 +50,57 @@ def main():
         layers = net_caffe.layers
     else:
         layers = parse.parse_caffemodel(args.caffe_model)
-    
+
     arg_shapes, output_shapes, aux_shapes = prob.infer_shape(data=tuple(input_dim))
     arg_names = prob.list_arguments()
     arg_shape_dic = dict(zip(arg_names, arg_shapes))
     arg_params = {}
-    
+
     iter = ''
     if caffe_flag:
         iter = get_caffe_iter(layer_names, layers)
     else:
         iter = get_iter(layers)
     first_conv = True
-    
+
     for layer_name, layer_type, layer_blobs in iter:
         if layer_type == 'Convolution' or layer_type == 'InnerProduct' or layer_type == 4 or layer_type == 14 \
                 or layer_type == 'PReLU':
             if layer_type == 'PReLU':
-                assert(len(layer_blobs) == 1)
+                assert (len(layer_blobs) == 1)
                 wmat = layer_blobs[0].data
                 weight_name = layer_name + '_gamma'
                 arg_params[weight_name] = mx.nd.zeros(wmat.shape)
                 arg_params[weight_name][:] = wmat
                 continue
-            assert(len(layer_blobs) == 2)
+            assert (len(layer_blobs) == 2)
             wmat_dim = []
             if getattr(layer_blobs[0].shape, 'dim', None) is not None:
                 if len(layer_blobs[0].shape.dim) > 0:
                     wmat_dim = layer_blobs[0].shape.dim
                 else:
-                    wmat_dim = [layer_blobs[0].num, layer_blobs[0].channels, layer_blobs[0].height, layer_blobs[0].width]
+                    wmat_dim = [layer_blobs[0].num, layer_blobs[0].channels, layer_blobs[0].height,
+                                layer_blobs[0].width]
             else:
                 wmat_dim = list(layer_blobs[0].shape)
             wmat = np.array(layer_blobs[0].data).reshape(wmat_dim)
             bias = np.array(layer_blobs[1].data)
             channels = wmat_dim[1]
-            if channels == 3 or channels == 4: # RGB or RGBA
+            if channels == 3 or channels == 4:  # RGB or RGBA
                 if first_conv:
-                    print 'Swapping BGR of caffe into RGB in mxnet'
+                    print('Swapping BGR of caffe into RGB in mxnet')
                     wmat[:, [0, 2], :, :] = wmat[:, [2, 0], :, :]
 
-            assert(wmat.flags['C_CONTIGUOUS'] is True)
-            assert(bias.flags['C_CONTIGUOUS'] is True)
-            print 'converting layer {0}, wmat shape = {1}, bias shape = {2}'.format(layer_name, wmat.shape, bias.shape)
+            assert (wmat.flags['C_CONTIGUOUS'] is True)
+            assert (bias.flags['C_CONTIGUOUS'] is True)
+            print('converting layer {0}, wmat shape = {1}, bias shape = {2}'.format(layer_name, wmat.shape, bias.shape))
             wmat = wmat.reshape((wmat.shape[0], -1))
             bias = bias.reshape((bias.shape[0], 1))
             weight_name = layer_name + "_weight"
             bias_name = layer_name + "_bias"
-            
+
             if weight_name not in arg_shape_dic:
-                print weight_name + ' not found in arg_shape_dic.'
+                print(weight_name + ' not found in arg_shape_dic.')
                 continue
             wmat = wmat.reshape(arg_shape_dic[weight_name])
             arg_params[weight_name] = mx.nd.zeros(wmat.shape)
@@ -108,11 +113,12 @@ def main():
             if first_conv and (layer_type == 'Convolution' or layer_type == 4):
                 first_conv = False
 
-    model = mx.model.FeedForward(ctx=mx.cpu(), symbol=prob,
-            arg_params=arg_params, aux_params={}, num_epoch=1,
-            learning_rate=0.05, momentum=0.9, wd=0.0001)
+    model = mx.mod.Module(symbol=prob, label_names=['prob_label', ])
+    model.bind(data_shapes=[('data', tuple(input_dim))])
+    model.init_params(arg_params=arg_params, aux_params={})
+
+    model.save_checkpoint(args.save_model_name, 1)
 
-    model.save(args.save_model_name)
 
 if __name__ == '__main__':
     main()

tools/caffe_converter/convert_symbol.py

@@ -1,3 +1,4 @@
+from __future__ import print_function
 from google.protobuf import text_format
 import argparse
 import re
@@ -11,23 +12,26 @@
     caffe_flag = False
     import caffe_parse.caffe_pb2
 
-def readProtoSolverFile(filepath):
+
+def read_proto_solver_file(file_path):
     solver_config = ''
     if caffe_flag:
         solver_config = caffe.proto.caffe_pb2.NetParameter()
     else:
         solver_config = caffe_parse.caffe_pb2.NetParameter()
-    return readProtoFile(filepath, solver_config)
+    return read_proto_file(file_path, solver_config)
+
 
-def readProtoFile(filepath, parser_object):
-    file = open(filepath, "r")
+def read_proto_file(file_path, parser_object):
+    file = open(file_path, "r")
     if not file:
-        raise self.ProcessException("ERROR (" + filepath + ")!")
+        raise Exception("ERROR (" + file_path + ")!")
     text_format.Merge(str(file.read()), parser_object)
     file.close()
     return parser_object
 
-def convParamToString(param):
+
+def conv_param_to_string(param):
     pad = 0
     if isinstance(param.pad, int):
         pad = param.pad
@@ -49,16 +53,16 @@ def convParamToString(param):
     else:
         dilate = 1 if len(param.dilation) == 0 else param.dilation[0]
     # convert to string except for dilation
-    param_string = "num_filter=%d, pad=(%d,%d), kernel=(%d,%d), stride=(%d,%d), no_bias=%s" %\
-        (param.num_output, pad, pad, kernel_size,\
-        kernel_size, stride, stride, not param.bias_term)
+    param_string = "num_filter=%d, pad=(%d,%d), kernel=(%d,%d), stride=(%d,%d), no_bias=%s" % \
+                   (param.num_output, pad, pad, kernel_size, kernel_size, stride, stride, not param.bias_term)
     # deal with dilation. Won't be in deconvolution
     if dilate > 1:
         param_string += ", dilate=(%d, %d)" % (dilate, dilate)
     return param_string
 
+
 def proto2script(proto_file):
-    proto = readProtoSolverFile(proto_file)
+    proto = read_proto_solver_file(proto_file)
     connection = dict()
     symbols = dict()
     top = dict()
@@ -70,52 +74,51 @@ def proto2script(proto_file):
     elif len(proto.layers):
         layer = proto.layers
     else:
-        raise Exception('Invalid proto file.')   
-    # Get input size to network
-    input_dim = [1, 3, 224, 224] # default
+        raise Exception('Invalid proto file.')
+        # Get input size to network
+    input_dim = [1, 3, 224, 224]  # default
     if len(proto.input_dim) > 0:
         input_dim = proto.input_dim
-    elif len(proto.input_shape) > 0: 
+    elif len(proto.input_shape) > 0:
         input_dim = proto.input_shape[0].dim
-    elif (layer[0].type == "Input"):
+    elif layer[0].type == "Input":
         input_dim = layer[0].input_param.shape._values[0].dim
         layer.pop(0)
     else:
-        raise Exception('Invalid proto file.')   
+        raise Exception('Invalid proto file.')
 
-    # We assume the first bottom blob of first layer is the output from data layer
+        # We assume the first bottom blob of first layer is the output from data layer
     input_name = layer[0].bottom[0]
     output_name = ""
-    mapping = {input_name : 'data'}
-    need_flatten = {input_name : False}
+    mapping = {input_name: 'data'}
+    need_flatten = {input_name: False}
     for i in range(len(layer)):
         type_string = ''
         param_string = ''
         name = re.sub('[-/]', '_', layer[i].name)
         if layer[i].type == 'Convolution' or layer[i].type == 4:
             type_string = 'mx.symbol.Convolution'
-            param_string = convParamToString(layer[i].convolution_param)
+            param_string = conv_param_to_string(layer[i].convolution_param)
             need_flatten[name] = True
         if layer[i].type == 'Deconvolution' or layer[i].type == 39:
             type_string = 'mx.symbol.Deconvolution'
-            param_string = convParamToString(layer[i].convolution_param)
+            param_string = conv_param_to_string(layer[i].convolution_param)
             need_flatten[name] = True
         if layer[i].type == 'Pooling' or layer[i].type == 17:
             type_string = 'mx.symbol.Pooling'
             param = layer[i].pooling_param
             param_string = ''
             param_string += "pooling_convention='full', "
-            if param.global_pooling == True:
+            if param.global_pooling:
                 # there must be a param `kernel` in a pooling layer
                 param_string += "global_pool=True, kernel=(1,1)"
             else:
-                param_string += "pad=(%d,%d), kernel=(%d,%d), stride=(%d,%d)" %\
-                    (param.pad, param.pad, param.kernel_size,\
-                    param.kernel_size, param.stride, param.stride)
+                param_string += "pad=(%d,%d), kernel=(%d,%d), stride=(%d,%d)" % \
+                                (param.pad, param.pad, param.kernel_size, param.kernel_size, param.stride, param.stride)
             if param.pool == 0:
-                param_string = param_string + ", pool_type='max'"
+                param_string += ", pool_type='max'"
             elif param.pool == 1:
-                param_string = param_string + ", pool_type='avg'"
+                param_string += ", pool_type='avg'"
             else:
                 raise Exception("Unknown Pooling Method!")
             need_flatten[name] = True
@@ -134,8 +137,8 @@ def proto2script(proto_file):
         if layer[i].type == 'LRN' or layer[i].type == 15:
             type_string = 'mx.symbol.LRN'
             param = layer[i].lrn_param
-            param_string = "alpha=%f, beta=%f, knorm=%f, nsize=%d" %\
-                (param.alpha, param.beta, param.k, param.local_size)
+            param_string = "alpha=%f, beta=%f, knorm=%f, nsize=%d" % \
+                           (param.alpha, param.beta, param.k, param.local_size)
             need_flatten[name] = True
         if layer[i].type == 'InnerProduct' or layer[i].type == 14:
             type_string = 'mx.symbol.FullyConnected'
@@ -179,31 +182,35 @@ def proto2script(proto_file):
             if len(bottom) == 1:
                 if need_flatten[mapping[bottom[0]]] and type_string == 'mx.symbol.FullyConnected':
                     flatten_name = "flatten_%d" % flatten_count
-                    symbol_string += "%s=mx.symbol.Flatten(name='%s', data=%s)\n" %\
-                        (flatten_name, flatten_name, mapping[bottom[0]])
+                    symbol_string += "%s=mx.symbol.Flatten(name='%s', data=%s)\n" % \
+                                     (flatten_name, flatten_name, mapping[bottom[0]])
                     flatten_count += 1
                     need_flatten[flatten_name] = False
                     bottom[0] = flatten_name
                     mapping[bottom[0]] = bottom[0]
-                symbol_string += "%s = %s(name='%s', data=%s %s)\n" %\
-                    (name, type_string, name, mapping[bottom[0]], param_string)
+                symbol_string += "%s = %s(name='%s', data=%s %s)\n" % \
+                                 (name, type_string, name, mapping[bottom[0]], param_string)
             else:
-                symbol_string += "%s = %s(name='%s', *[%s] %s)\n" %\
-                    (name, type_string, name, ','.join([mapping[x] for x in bottom]), param_string)
+                symbol_string += "%s = %s(name='%s', *[%s] %s)\n" % \
+                                 (name, type_string, name, ','.join([mapping[x] for x in bottom]), param_string)
         for j in range(len(layer[i].top)):
             mapping[layer[i].top[j]] = name
         output_name = name
     return symbol_string, output_name, input_dim
 
+
 def proto2symbol(proto_file):
     sym, output_name, input_dim = proto2script(proto_file)
     sym = "import mxnet as mx\n" \
-            + "data = mx.symbol.Variable(name='data')\n" \
-            + sym
+          + "data = mx.symbol.Variable(name='data')\n" \
+          + sym
     exec(sym)
-    exec("ret = " + output_name)
+    _locals = locals()
+    exec("ret = " + output_name, globals(), _locals)
+    ret = _locals['ret']
     return ret, input_dim
 
+
 def main():
     symbol_string, output_name, input_dim = proto2script(sys.argv[1])
     if len(sys.argv) > 2:
@@ -212,5 +219,6 @@ def main():
     else:
         print(symbol_string)
 
+
 if __name__ == '__main__':
     main()

tools/caffe_converter/mean_image.py

@@ -10,11 +10,12 @@
     caffe_flag = False
     import caffe_parse.caffe_pb2
 
-def protoBlobFileToND(protofile):
+
+def protoBlobFileToND(proto_file):
     data = ''
-    file = open(protofile, "r")
+    file = open(proto_file, "r")
     if not file:
-        raise self.ProcessException("ERROR (" + protofile + ")!")
+        raise Exception("ERROR (" + proto_file + ")!")
     data = file.read()
     file.close()
 
@@ -34,6 +35,7 @@ def protoBlobFileToND(protofile):
     img_mean_np[2] = img_mean_np2[0]
     return mx.nd.array(img_mean_np)
 
+
 def main():
     parser = argparse.ArgumentParser(description='Caffe prototxt to mxnet model parameter converter.\
                     Note that only basic functions are implemented. You are welcomed to contribute to this file.')
@@ -45,4 +47,4 @@ def main():
 
 
 if __name__ == '__main__':
-    main()
+    main()