pytorch转caffe2模型,方法1亲测可用

原创

pytorch转caffe2模型

2019-05-30 17:19:16 酉意铭 阅读数 932 收藏 更多

分类专栏： 模型转换

版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。

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

            </div>
                                                <!--一个博主专栏付费入口-->
         
         <!--一个博主专栏付费入口结束-->
        <link rel="stylesheet" href="https://csdnimg.cn/release/phoenix/template/css/ck_htmledit_views-4a3473df85.css">
                                    <link rel="stylesheet" href="https://csdnimg.cn/release/phoenix/template/css/ck_htmledit_views-4a3473df85.css">
            <div class="htmledit_views" id="content_views">
                                        <p>&nbsp; &nbsp; 项目需要用c++调用caffe2接口进行工程化部署，但是模型是用pytorch实现并训练。这就需要把pytorch预训练模型转化为可供c++调用的模型。</p>

 

    目前的问题是：训练环境和工程部署环境是不同的。利用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
 
 

 

文章最后发布于: 2019-05-30 17:19:16