| """Convert caffe prototxt to symbol |
| """ |
| from __future__ import print_function |
| import argparse |
| import re |
| import caffe_parser |
| |
| def _get_input(proto): |
| """Get input size |
| """ |
| layer = caffe_parser.get_layers(proto) |
| if len(proto.input_dim) > 0: |
| input_dim = proto.input_dim |
| elif len(proto.input_shape) > 0: |
| input_dim = proto.input_shape[0].dim |
| elif layer[0].type == "Input": |
| input_dim = layer[0].input_param.shape[0].dim |
| layer.pop(0) |
| else: |
| raise ValueError('Cannot find input size') |
| |
| assert layer[0].type != "Input", 'only support single input' |
| # We assume the first bottom blob of first layer is the output from data layer |
| input_name = layer[0].bottom[0] |
| return input_name, input_dim, layer |
| |
| def _convert_conv_param(param): |
| """ |
| Convert convolution layer parameter from Caffe to MXNet |
| """ |
| param_string = "num_filter=%d" % param.num_output |
| |
| pad_w = 0 |
| pad_h = 0 |
| if isinstance(param.pad, int): |
| pad = param.pad |
| param_string += ", pad=(%d, %d)" % (pad, pad) |
| else: |
| if len(param.pad) > 0: |
| pad = param.pad[0] |
| param_string += ", pad=(%d, %d)" % (pad, pad) |
| else: |
| if isinstance(param.pad_w, int): |
| pad_w = param.pad_w |
| if isinstance(param.pad_h, int): |
| pad_h = param.pad_h |
| param_string += ", pad=(%d, %d)" % (pad_h, pad_w) |
| |
| if isinstance(param.kernel_size, int): |
| kernel_size = param.kernel_size |
| param_string += ", kernel=(%d,%d)" % (kernel_size, kernel_size) |
| else: |
| if len(param.kernel_size) > 0: |
| kernel_size = param.kernel_size[0] |
| param_string += ", kernel=(%d,%d)" % (kernel_size, kernel_size) |
| else: |
| assert isinstance(param.kernel_w, int) |
| kernel_w = param.kernel_w |
| assert isinstance(param.kernel_h, int) |
| kernel_h = param.kernel_h |
| param_string += ", kernel=(%d,%d)" % (kernel_h, kernel_w) |
| |
| stride = 1 |
| if isinstance(param.stride, int): |
| stride = param.stride |
| else: |
| stride = 1 if len(param.stride) == 0 else param.stride[0] |
| |
| param_string += ", stride=(%d,%d)" % (stride, stride) |
| |
| dilate = 1 |
| if hasattr(param, 'dilation'): |
| if isinstance(param.dilation, int): |
| dilate = param.dilation |
| else: |
| dilate = 1 if len(param.dilation) == 0 else param.dilation[0] |
| |
| param_string += ", no_bias=%s" % (not param.bias_term) |
| |
| # deal with dilation. Won't be in deconvolution |
| if dilate > 1: |
| param_string += ", dilate=(%d, %d)" % (dilate, dilate) |
| |
| if isinstance(param.group, int): |
| if param.group != 1: |
| param_string += ", num_group=%d" % param.group |
| |
| return param_string |
| |
| def _convert_pooling_param(param): |
| """Convert the pooling layer parameter |
| """ |
| param_string = "pooling_convention='full', " |
| if param.global_pooling: |
| 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) |
| if param.pool == 0: |
| param_string += ", pool_type='max'" |
| elif param.pool == 1: |
| param_string += ", pool_type='avg'" |
| else: |
| raise ValueError("Unknown Pooling Method!") |
| return param_string |
| |
| def _parse_proto(prototxt_fname): |
| """Parse Caffe prototxt into symbol string |
| """ |
| proto = caffe_parser.read_prototxt(prototxt_fname) |
| |
| # process data layer |
| input_name, input_dim, layers = _get_input(proto) |
| # only support single input, so always use `data` as the input data |
| mapping = {input_name: 'data'} |
| need_flatten = {input_name: False} |
| symbol_string = "import mxnet as mx\ndata = mx.symbol.Variable(name='data')\n" |
| |
| flatten_count = 0 |
| output_name = "" |
| prev_name = None |
| |
| # convert reset layers one by one |
| for i, layer in enumerate(layers): |
| type_string = '' |
| param_string = '' |
| skip_layer = False |
| name = re.sub('[-/]', '_', layer.name) |
| if layer.type == 'Convolution' or layer.type == 4: |
| type_string = 'mx.symbol.Convolution' |
| param_string = _convert_conv_param(layer.convolution_param) |
| need_flatten[name] = True |
| if layer.type == 'Deconvolution' or layer.type == 39: |
| type_string = 'mx.symbol.Deconvolution' |
| param_string = _convert_conv_param(layer.convolution_param) |
| need_flatten[name] = True |
| if layer.type == 'Pooling' or layer.type == 17: |
| type_string = 'mx.symbol.Pooling' |
| param_string = _convert_pooling_param(layer.pooling_param) |
| need_flatten[name] = True |
| if layer.type == 'ReLU' or layer.type == 18: |
| type_string = 'mx.symbol.Activation' |
| param_string = "act_type='relu'" |
| need_flatten[name] = need_flatten[mapping[layer.bottom[0]]] |
| if layer.type == 'TanH' or layer.type == 23: |
| type_string = 'mx.symbol.Activation' |
| param_string = "act_type='tanh'" |
| need_flatten[name] = need_flatten[mapping[layer.bottom[0]]] |
| if layer.type == 'Sigmoid' or layer.type == 19: |
| type_string = 'mx.symbol.Activation' |
| param_string = "act_type='sigmoid'" |
| need_flatten[name] = need_flatten[mapping[layer.bottom[0]]] |
| if layer.type == 'LRN' or layer.type == 15: |
| type_string = 'mx.symbol.LRN' |
| param = layer.lrn_param |
| 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.type == 'InnerProduct' or layer.type == 14: |
| type_string = 'mx.symbol.FullyConnected' |
| param = layer.inner_product_param |
| param_string = "num_hidden=%d, no_bias=%s" % ( |
| param.num_output, not param.bias_term) |
| need_flatten[name] = False |
| if layer.type == 'Dropout' or layer.type == 6: |
| type_string = 'mx.symbol.Dropout' |
| param = layer.dropout_param |
| param_string = "p=%f" % param.dropout_ratio |
| need_flatten[name] = need_flatten[mapping[layer.bottom[0]]] |
| if layer.type == 'Softmax' or layer.type == 20: |
| type_string = 'mx.symbol.SoftmaxOutput' |
| if layer.type == 'Flatten' or layer.type == 8: |
| type_string = 'mx.symbol.Flatten' |
| need_flatten[name] = False |
| if layer.type == 'Split' or layer.type == 22: |
| type_string = 'split' # will process later |
| if layer.type == 'Concat' or layer.type == 3: |
| type_string = 'mx.symbol.Concat' |
| need_flatten[name] = True |
| if layer.type == 'Crop': |
| type_string = 'mx.symbol.Crop' |
| need_flatten[name] = True |
| param_string = 'center_crop=True' |
| if layer.type == 'BatchNorm': |
| type_string = 'mx.symbol.BatchNorm' |
| param = layer.batch_norm_param |
| # CuDNN requires eps to be greater than 1e-05 |
| # We compensate for this change in convert_model |
| epsilon = param.eps |
| if (epsilon <= 1e-05): |
| epsilon = 1e-04 |
| # if next layer is scale, don't fix gamma |
| fix_gamma = layers[i+1].type != 'Scale' |
| param_string = 'use_global_stats=%s, fix_gamma=%s, eps=%f' % ( |
| param.use_global_stats, fix_gamma, epsilon) |
| need_flatten[name] = need_flatten[mapping[layer.bottom[0]]] |
| if layer.type == 'Scale': |
| assert layers[i-1].type == 'BatchNorm' |
| need_flatten[name] = need_flatten[mapping[layer.bottom[0]]] |
| skip_layer = True |
| prev_name = re.sub('[-/]', '_', layers[i-1].name) |
| if layer.type == 'PReLU': |
| type_string = 'mx.symbol.LeakyReLU' |
| param = layer.prelu_param |
| param_string = "act_type='prelu', slope=%f" % param.filler.value |
| need_flatten[name] = need_flatten[mapping[layer.bottom[0]]] |
| if layer.type == 'Eltwise': |
| type_string = 'mx.symbol.broadcast_add' |
| param_string = "" |
| need_flatten[name] = False |
| if layer.type == 'Reshape': |
| type_string = 'mx.symbol.Reshape' |
| need_flatten[name] = False |
| param = layer.reshape_param |
| param_string = "shape=(%s)" % (','.join(param.shape.dim),) |
| if layer.type == 'AbsVal': |
| type_string = 'mx.symbol.abs' |
| need_flatten[name] = need_flatten[mapping[layer.bottom[0]]] |
| |
| if skip_layer: |
| assert len(layer.bottom) == 1 |
| symbol_string += "%s = %s\n" % (name, prev_name) |
| elif type_string == '': |
| raise ValueError('Unknown layer %s!' % layer.type) |
| elif type_string != 'split': |
| bottom = layer.bottom |
| if param_string != "": |
| param_string = ", " + param_string |
| 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]]) |
| 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) |
| else: |
| 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.top)): |
| mapping[layer.top[j]] = name |
| output_name = name |
| return symbol_string, output_name, input_dim |
| |
| def convert_symbol(prototxt_fname): |
| """Convert caffe model definition into Symbol |
| |
| Parameters |
| ---------- |
| prototxt_fname : str |
| Filename of the prototxt file |
| |
| Returns |
| ------- |
| Symbol |
| Converted Symbol |
| tuple |
| Input shape |
| """ |
| sym, output_name, input_dim = _parse_proto(prototxt_fname) |
| exec(sym) # pylint: disable=exec-used |
| _locals = locals() |
| exec("ret = " + output_name, globals(), _locals) # pylint: disable=exec-used |
| ret = _locals['ret'] |
| return ret, input_dim |
| |
| def main(): |
| parser = argparse.ArgumentParser( |
| description='Convert caffe prototxt into Symbol') |
| parser.add_argument('prototxt', help='The prototxt filename') |
| parser.add_argument('output', help='filename for the output json file') |
| args = parser.parse_args() |
| |
| sym, _ = convert_symbol(args.prototxt) |
| sym.save(args.output) |
| |
| if __name__ == '__main__': |
| main() |
