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