caffe2 mdl文件转init_net.pb, predict_net.pb

前言

caffe2中，完整的训练例子就train_resnet50.py, 其模型和snapshot都保持在同一个文件中，不像caffe 一样有.caffemodel 和 .snapshot 文件（其实是冗余的）。训练完之后发现*.mdl文件不能直接使用，需要转化为两个pb文件，然而没有说怎么转换，也有很多人在问这个issue，然而并没有人说出正确的方法。

解决方法

通过看代码（好像也没有别的方法了），发现了可以通过

from caffe2.python.predictor.mobile_exporter import Export

这个文件函数来进行转换。高高兴兴的转换了，一测试，又报错说啥*-rm *_riv没有定义啥的，也就是说关于BN的运行均值和方差没有存入到网络中，需要手动指定一下（没有BN的网络没有这个问题），简单的包装一下，如下：

from caffe2.python import core, model_helper, net_drawer, workspace, visualize, brew
from caffe2.python.predictor.mobile_exporter import Export
def export(workspace, net, params, init_net_name="init_net.pb", predict_net_name="predict_net.pb" ):
    extra_params = []
    extra_blobs = []
    for blob in workspace.Blobs():
    |   name = str(blob)
    |   if name.endswith("_rm") or name.endswith("_riv"):
    |   |   extra_params.append(name)
    |   |   extra_blobs.append(workspace.FetchBlob(name))
    for name, blob in zip(extra_params, extra_blobs):
    |   workspace.FeedBlob(name, blob)
    |   params.append(name)
    init_net, predict_net = Export(workspace, net, params)
    with open(init_net_name, 'wb') as f:
    |   f.write(init_net.SerializeToString())
    with open(predict_net_name, 'wb') as f:
    |   f.write(predict_net.SerializeToString())

这样，只需要在需要的保存的时候，调用

export(workspace, deploy_model.net, deploy_model.params)

就可以了。