前言

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)

就可以了。

總結

以上是生活随笔為你收集整理的caffe2 mdl文件转init_net.pb, predict_net.pb的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。