caffe2之caffe_translator

来源:互联网 发布:java velocity 编辑:程序博客网 时间:2024/06/03 15:51

前言

博主最近在做算法移植,也就是把caffe上面训练好的模型移植到caffe2上面,踩过了不少的坑,趁着周末总结一下经验。本文主要分为三个模块1),如何使用caffe2的python目录下的caffe_translator.py工具来实现常规模型的转化;2),部分代码的解读;3,动手实现层转化,caffe里面一些层在caffe2里面并没有实现,所以这时候需要自己添加层的转化,(讲道理并不难,哈哈)

1 useage

参照官网可以使用的这样的命令来实现模型的转化,如下所示:

python -m caffe2.python.caffe_translator deploy.prototxt  pretrained.caffemodel

上面的命令就是转化不需要考虑的地方,当然博主比较喜欢用下面的命令。

python caffe_translator.py  deploy.prototxt pretrain.caffemodel [xxx.pb]    [xxx.pb] 

后面两个可以不输入,或者取你喜欢的名字,具体原因可以看看源码就是了解了。

2 部分代码解读

本部分写的主要目的是对caffe_translator有了整体上的认识,从而有助于方面以后找bug(不要以为官网上的代码就没问题了,哈哈),同时也会对层的添加有了清楚的认识。

2.1 整体的思路

将输入的.prototxt 和 .caffemodel根据提供的caffe.proto,将op和params解析出来,然后根据caffe2.proto生成对应的文件。其中很多模块的特征都是从proto里面来的。

2.2 解读部分代码

2.2.1 main函数
2.2.2 TranslateMode

##########################################################function:translatemodel#input:#1,caffe_net(read from proto);#2,pretrainde_net(read #from caffemodel)#3,is_test=False,#4,net_state(none)#output:#net(for caffe2),net_params(caffe2)#commomt:代码的核心,目的是用于模型的转化。#值得注意,net和net_params 两者的属性;#net = caffe2_pb2.NetDef();此处定义了net的属性,网络的#input,output和op等都是包含在里面,类似与caffe的prototxt。#NetDef是个类,它的定义在caffe2.proto里面,并且通过#protobuffer生成的。具体可以看下面的代码,##########################################################message NetDef {#  optional string name = 1; // the network's name#  repeated OperatorDef op = 2;#  optional string type = 3;#  optional int32 num_workers = 4 [deprecated=true];#  optional DeviceOption device_option = 5;#  repeated Argument arg = 6;#  repeated string external_input = 7;#  repeated string external_output = 8;#}########################################################## net_params = caffe2_pb2.TensorProtos()#可以理解成caffe里面的caffemodel,当然从功能上来说,可以存放#tensor的目标,存放数据的地方。#########################################################    def TranslateModel(        cls,        caffe_net,        pretrained_net,        is_test=False,        net_state=None,    ):        net_state = caffe_pb2.NetState() if net_state is None else net_state        net = caffe2_pb2.NetDef()              net.name = caffe_net.name        net_params = caffe2_pb2.TensorProtos()        if len(caffe_net.layers) > 0:            raise ValueError(                'I think something is wrong. This translation script '                'only accepts new style layers that are stored in the '                'layer field.'            )#遍历caffe_net里面的所有层,检验里面网络里面是否有重名的层。检验#之后开始转化。        for layer in caffe_net.layer:            if not _ShouldInclude(net_state, layer):                log.info('Current net state does not need layer {}'                            .format(layer.name))                continue            log.info('Translate layer {}'.format(layer.name))            # Get pretrained one            pretrained_layers = (                [l for l in pretrained_net.layer                 if l.name == layer.name] + [l                                             for l in pretrained_net.layers                                             if l.name == layer.name]            )            if len(pretrained_layers) > 1:                raise ValueError(                    'huh? more than one pretrained layer of one name?')            elif len(pretrained_layers) == 1:            #当网络的定义为这名字额网络只有一层,则将模型的参数            #序列化,生成列表                pretrained_blobs = [                    utils.CaffeBlobToNumpyArray(blob)                    for blob in pretrained_layers[0].blobs                ]            else:                # No pretrained layer for the given layer name. We'll just pass                # no parameter blobs.                # print 'No pretrained layer for layer', layer.name                pretrained_blobs = []            operators, params = cls.TranslateLayer(                layer, pretrained_blobs, is_test)            net.op.extend(operators)            net_params.protos.extend(params)        return net, net_params

2.2.3 TranslateLayer
此处在做模型的参数和定义的转化工作。这个转化工具里面一些common layer 的转化已经写好,但是有些层就没有写了,例如priorbox layer 等等。如果添加新的层的话,就在这里改一下就可以了。

##########################################################function:translatelayer#inpute:1,layer;2,corresponding  layer_params;#output:caffe_ops(op),params(参数)  ()for  one  layer )##########################################################    @classmethod    def TranslateLayer(cls, layer, pretrained_blobs, is_test):        try:            caffe_ops, params = cls.registry_[layer.type](                layer, pretrained_blobs, is_test)    #caffe_ops, params = func(layer, pretrained_blobs, is_test)        except KeyError:            raise KeyError('No translator registered for layer: %s yet.' %                           str(layer))        if caffe_ops is None:            caffe_ops = []        if type(caffe_ops) is not list:            caffe_ops = [caffe_ops]        return caffe_ops, params
 caffe_ops, params = cls.registry_[layer.type](                layer, pretrained_blobs, is_test) @TranslatorRegistry.Register("LRN")##########################################################此处就是为了注册层的功能,@是python里面的装饰器,可以这么理解,名词调用相当于注册这个层。 TranslateLRN(layer, pretrained_blobs, is_test):    caffe_op = BaseTranslate(layer, "LRN")    #caffe2_op = caffe2_pb2.OperatorDef()##########################################################message OperatorDef {#  repeated string input = 1; // the name of the input blobs#  repeated string output = 2; // the name of output top blobs#  optional string name = 3; // the operator name. This is optional.#  optional string type = 4;#  repeated Argument arg = 5;#  optional DeviceOption device_option = 6;#  optional string engine = 7;#  repeated string control_input = 8;#  optional bool is_gradient_op = 9 [default = false];#}#################################################    caffe_op.output.extend(['_' + caffe_op.output[0] + '_scale'])    param = layer.lrn_param    #获取caffe_net网络层里面的超参数    if param.norm_region != caffe_pb2.LRNParameter.ACROSS_CHANNELS:        raise ValueError(            "Does not support norm region other than across channels.")    #此处可以链接到MakeArgument,在op里面注册到Argument下面。    AddArgument(caffe_op, "size", int(param.local_size))    AddArgument(caffe_op, "alpha", float(param.alpha))    AddArgument(caffe_op, "beta", float(param.beta))    AddArgument(caffe_op, "bias", float(param.k))    AddArgument(caffe_op, "order", "NCHW")    return caffe_op, []    #返回op的对象,和模型的参数(这里就没有参数返回了。)

3 算法移植

例如上面的层里面没有priorbox_layer

3.1 查看caffe.proto里面关于的描述

message PriorBoxParameter {  enum CodeType {    CORNER = 1;    CENTER_SIZE = 2;    CORNER_SIZE = 3;  }  repeated float min_size = 1;  repeated float max_size = 2;  repeated float aspect_ratio = 3;  optional bool flip = 4 [default = true];  optional bool clip = 5 [default = false];  repeated float variance = 6;  optional uint32 img_size = 7;  optional uint32 img_h = 8;  optional uint32 img_w = 9;  optional float step = 10;  optional float step_h = 11;  optional float step_w = 12;  optional float offset = 13 [default = 0.5];}

3.2 注册层

@TranslatorRegistry.Register("PriorBox")

3.3 取个名字

caffe_op = BaseTranslate(layer, "PriorBox")

3.4 获取layer_param

param = layer.prior_box_param

3.5 添加到op里面

AddArgument(caffe_op, "min_sizes", [param.min_size[i] for i in range(len(param.min_size))])AddArgument(caffe_op, "max_sizes", [param.max_size[i] for i in range(len(param.max_size))])AddArgument(caffe_op, "aspect_ratios", [param.aspect_ratio[i] for i in range(len(param.aspect_ratio))])AddArgument(caffe_op, "flip", param.flip)AddArgument(caffe_op, "clip", param.clip)AddArgument(caffe_op, "variance", [param.variance[i] for i in range(len(param.variance))])AddArgument(caffe_op, "img_size", param.img_size)AddArgument(caffe_op, "img_w", param.img_w)AddArgument(caffe_op, "img_h", param.img_h)AddArgument(caffe_op, "step", param.step)AddArgument(caffe_op, "step_h", param.step_h)AddArgument(caffe_op, "step_w", param.step_w)AddArgument(caffe_op, "offset", param.offset)AddArgument(caffe_op, "order", "NCHW")