Caffe模型移植到MXNet

DSP

Caffe模型移植到MXNet

2019-07-13 17:38发布生成海报

站内文章 / DSP

10886 0

使用caffe的一大好处是有很多的预训练模型，你可以从caffe的model zoo去下载这些模型。那么怎样把caffe的模型转到MXNet中呢？一种最简单也是最有效的方法就是把caffe的模型加载出来，然后对照着模型参数，逐个复制到MXNet对应的模型参数中。这种方法简单有效，但是也是工作量比较大的一种方法。其实MXNet提供了相应的转换工具帮助我们完成这一流程，本文记录一下这种方法。

下载caffe模型

我需要在MXNet上使用人脸特征点检测，在caffe model zoo中我找到了一个叫VanillaCNN的模型：

https://github.com/ishay2b/VanillaCNN

caffe的模型有两个文件vanillaCNN.caffemodel和vanilla_deploy.prototxt，把这两个文件下载下来。

转换工具

这个转换工具其实就在MXNet的源代码中：

https://github.com/dmlc/mxnet/tree/master/tools

进入mxnet/tools/caffe_converter目录下，convert_model.py就是这里使用的转换工具。根据页面提示，需要安装protobuf的相关工具：

sudo apt-get install protobuf-compiler
sudo pip install protobuf

安装好了以后，使用这个工具很简单： python convert_model.py vanilla_deploy.prototxt vanilla.caffemodel vanilla convert_model.py第一个参数是prototxt，第二个参数是caffemodel，第三个参数是要生成的mxnet的模型名称。不出意外，这一步会出错: Exception: Unknown Layer AbsVal! 原来是caffe中的AbsVal层无法识别。

Hack

这么一来，只能自己Hack了。观察一下这个目录中的源代码，我们发现convert_symbol.py中的代码很可疑，打开看一下：

        if layer[i].type == 'ReLU' or layer[i].type == 18:
            type_string = 'mx.symbol.Activation'
            param_string = "act_type='relu'"
            need_flatten[name] = need_flatten[mapping[layer[i].bottom[0]]]

        if layer[i].type == 'TanH' or layer[i].type == 23:
            type_string = 'mx.symbol.Activation'
            param_string = "act_type='tanh'"
            need_flatten[name] = need_flatten[mapping[layer[i].bottom[0]]]

这些代码显然是将caffe模型中的符号和mxnet相对应。那么我们要hack的就是这个文件。

符号对应

现在出错的提示是 Exception: Unknown Layer AbsVal! 那么AbsVal是干什么的？打开caffe的文档：

http://caffe.berkeleyvision.org/tutorial/layers.html

找到AbsVal：

Absolute Value
Layer type: AbsVal
CPU implementation: ./src/caffe/layers/absval_layer.cpp
CUDA GPU implementation: ./src/caffe/layers/absval_layer.cu
Sample

layer {
      name: "layer"
      bottom: "in"
      top: "out"
      type: "AbsVal"
}
The AbsVal layer computes the output as abs(x) for each input element x.

原来AbsVal就是对输入求绝对值，那么对应的mxnet符号是什么呢？再打开mxnet的文档：

http://mxnet.io/api/python/symbol.html

仔细找找，找到了这个

http://mxnet.io/api/python/symbol.html#mxnet.symbol.abs

阅读一下：

mxnet.symbol.abs(*args, **kwargs)

Take absolute value of the src

Parameters: 
    src (Symbol) – Left symbolic input to the function
    name (string, optional.) – Name of the resulting symbol.
Returns:    
    symbol – The result symbol.
Return type:    
    Symbol

也是将输入求绝对值，那么可以确定caffe中的AbsVal对应于MXNet中的mx.symbol.abs。

问题解决

根据上面的分析，我们就在convert_symbol.py加上这样的一句看看行不行：

      if layer[i].type == 'AbsVal':
            type_string = 'mx.symbol.abs'
            need_flatten[name] = need_flatten[mapping[layer[i].bottom[0]]]

这句话的写法是模仿上面的ReLU的写法，和RELU相比，abs没有别的参数了,所以 param_string需要删去,而abs和ReLU的相似之处在于它们都是对输入的一对一的映射，所以在need_flatten这个参数上是相似的。再试一次，出现：

Swapping BGR of caffe into RGB in mxnet
converting layer Conv1, wmat shape = (16, 3, 5, 5), bias shape = (16,)
converting layer Conv2, wmat shape = (48, 16, 3, 3), bias shape = (48,)
converting layer Conv3, wmat shape = (64, 48, 3, 3), bias shape = (64,)
converting layer Conv4, wmat shape = (64, 64, 2, 2), bias shape = (64,)
converting layer Dense1, wmat shape = (100, 576), bias shape = (100,)
converting layer Dense2, wmat shape = (10, 100), bias shape = (10,)

可以看到的确成功了，同时在目录下生成了MXNet模型所需要的两个文件：

vanilla-0001.params
vanilla-symbol.json

由此，我们可以认为转换成功。但是我们还需要在实际中对比一下两个模型的输出是否一致，这里就不展开了，可以查看我放在GitHub的代码，去试验一下模型的精度。这里需要注意的是第一行 Swapping BGR of caffe into RGB in mxnet 在Caffe中，图像的输入是BGR格式的，但是在MXNet中图像是RGB格式，在输入的时候需要注意一下。

源代码

我把转换好的模型放在GitHub上了,欢迎star一下：

https://github.com/flyingzhao/mxnet_VanillaCNN

总结

MXNet在很多地方都需要自己Hack，上次编译Android动态链接库也是需要自己进行Hack，但是MXNet代码写的很模块化，Hack起来也很方便，顺便读一读别人的代码，何乐而不为？