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

————————————————

版权声明：本文为CSDN博主「酉意铭」的原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接及本声明。

原文链接：https://blog.csdn.net/weixin_40671425/article/details/90700908