pytorch轉caffe2模型

    項目需要用c++調用caffe2接口進行工程化部署，但是模型是用pytorch實現並訓練。這就需要把pytorch預訓練模型轉化爲可供c++調用的模型。

 

    目前的問題是：訓練環境和工程部署環境是不同的。利用python3.6+pytorch進行訓練，模型部署是python2.7+caffe2。

    總體流程：

        1，在訓練環境下也就是python3.6+pytorch下利用torch.onnx包把pth文件轉化爲onnx文件

         2，在python2.7+caffe2環境下把轉換好的onnx文件轉化爲pb文件供c++調用

 

安裝編譯caffe2參考官網：https://caffe2.ai/docs/getting-started.html?platform=windows&configuration=compile

 

一，pytorch -> onnx

import numpy as np
import torch
import torch.onnx
from models.faceboxes import FaceBoxes


# 加載定義的網絡模型，預測試階段相同
model = FaceBoxes(phase='test', size=None, num_classes=2)
# 加載預訓練模型
pretrained_model = "./pretrained/face_head_brain_our/model_face_brain_ourFinal_FaceBoxes.pth"

# 把預訓練模型參數加載進網絡，這裏採用GPU環境 也可以採用CPU
device = torch.cuda.current_device()
pretrained_dict = torch.load(pretrained_model, map_location=lambda storage, loc: storage.cuda(device))
model.load_state_dict(pretrained_dict, strict=False)

# 將訓練模式設置爲false, 因爲只需要網絡forward
model.train(False)

# 生成一個隨機張量用於前傳，這個張量可以是一個圖像張量也可以是一個隨機張量，數值不重要，只要size匹配即可
batch_size = 1
x = torch.randn(batch_size,3,1024,1024,requires_grad=True)
print(type(x))

# 導出模型 pytorch -> onnx
"""
這裏也可以用torch.onnx.export 詳細見： 
	https://github.com/pytorch/pytorch/blob/master/torch/onnx/utils.py
但最後都是調用 _export()函數
"""
torch_out = torch.onnx._export(
	model,						# 具有預訓練參數的模型
	x,							# 輸入張量，dumm data
	"faceboxes.onnx",			# 輸出文件保存的路徑
	export_params=True			# 保存模型內部訓練好的模型參數
	)

二，onnx -> caffe2

    方法1：

import os
import onnx
from caffe2.python.onnx.backend import Caffe2Backend
#也可以用onnx_caffe2中的caffebackend包，需要額外安裝onnx_caffe2，但是已經廢棄了
# from onnx_caffe2.backend import Caffe2Backend
import argparse
import caffe2.python.utils as putils
from caffe2.python import core, workspace
from caffe2.proto import caffe2_pb2

onnx_proto_file = 'faceboxes-gpu.onnx'
model_name = 'faceboxes'
model_dir = './transform_rst'
if not os.path.exists(model_dir):
	os.makedirs(model_dir)

onnx_model = onnx.load(onnx_proto_file)
init_net, predict_net = Caffe2Backend.onnx_graph_to_caffe2_net(onnx_model, device="CUDA:0")

with open(os.path.join(model_dir, model_name + "_init.pb"), "wb") as f:
	f.write(init_net.SerializeToString())
#with open(os.path.join(model_dir, model_name + "_init.pbtxt"), "w") as f:
#	f.write(str(init_net))
with open(os.path.join(model_dir, model_name + "_predict.pb"), "wb") as f:
	f.write(predict_net.SerializeToString())

#保存爲pbtxt文件的格式方便查看轉換的網絡結構
"""
# 如果只有pb也可以用pb文件轉化爲pbtxtx文件，這兩中文件格式可以互相轉換
pb和pbtxt互轉參考： https://github.com/dzung-hoang/caffe2-utils 
"""
with open(os.path.join(model_dir, model_name + "_predict.pbtxt"), "w") as f:
	f.write(str(predict_net))

    方法2：

import onnx
import caffe2.python.onnx.backend as backend
import numpy as np

#參考： https://pytorch.org/tutorials/advanced/super_resolution_with_caffe2.html
"""
此方法的轉換可以在移動設備上運行，但是轉換的效果和用Caffe2Backend包一樣的
from caffe2.python.onnx.backend import Caffe2Backend

"""
batch_size = 1

dummy_data = np.random.randn(1,3,1024,1024).astype(np.float32)

model = onnx.load("faceboxes-gpu.onnx")
onnx.checker.check_model(model)
prepared_backend = backend.prepare(model)
rep = backend.prepare(model,device="CUDA:0")
output = rep.run(dummy_data)

W = {model.graph.input[0].name: dummy_data}

c2_out = rep.run(W)[0]
print(c2_out)

#np.testing.assert_almost_equal(torch_out.data.cpu().numpy(), c2_out, decimal=3)
#print("Exported model has been executed on Caffe2 backend, and the result looks good!")

c2_workspace = rep.workspace
c2_model = rep.predict_net

from caffe2.python.predictor import mobile_exporter
init_net, predict_net = mobile_exporter.Export(c2_workspace, c2_model, c2_model.external_input)
with open('init_net.pb', "wb") as fopen:
    fopen.write(init_net.SerializeToString())
with open('predict_net.pb', "wb") as fopen:
    fopen.write(predict_net.SerializeToString())

    方法3：直接採用命令行

    

convert-onnx-to-caffe2 assets/squeezenet.onnx --output predict_net.pb --init-net-output init_net.pb