CenterPoint 部署攻略二（pth2onnx 基于mmengine）（动态多输入、多模型、校准验证）

本文链接：https://blog.youkuaiyun.com/Xiaobai_Zhao/article/details/132668042

简述:

系列二，主要就是 pth2onnx 的脚本及注释，本篇不涉及自定义算子，因此也是其他模型的通用模板！

针对 centerpoint 的构造(centerpoint_pillar02_second_secfpn_8xb4-cyclic-20e_nus-3d)，即：Det3DDataPreprocessor -> PillarFeatureNet -> PointPillarsScatter -> SECOND -> SECONDFPN -> CenterHead；

其中preprocessdata, pps不涉及参数的传递，构造onnx的时候把他们排除，所以是 PFN 和 SECOND...CenterHead 两个模型的游戏。

细节：

1. 加载模型：

import time
import torch
import onnx
import numpy as np
import onnxruntime as ort
from mmengine.config import Config
from mmdet3d.registry import MODELS
from onnxsim import simplify
from mmengine.model import BaseModule


cfg = Config.fromfile(cfg_path)
model = MODELS.build(cfg.model)
loaded_ckpt = torch.load(ckpt_path, map_location='cpu')
model.load_state_dict(loaded_ckpt['state_dict'], strict=True)
model.eval()
model = model.cuda()

“cfg_path” 换成你自己的config，load_state_dict 中 strict 的意思是你的 pth 保存的权重（名称）是否和 config 中搭建的模型保持严格一致；

注意：model.eval() 一定要加！

至此，得到了一个含权重的“总”模型。

按上述所言，需要提取出两个阶段的单独的模型，因为 PFN 直接通过 model.pts_voxel_encoder 就可以搭建（偷懒了），所以考虑第二阶段的模型构造：

class Stage2Model(BaseModule):
    def __init__(self, model):
        super(Stage2Model, self).__init__()
        self.backbone = model.pts_backbone
        self.neck = model.pts_neck
        self.head = model.pts_bbox_head

    def forward(self, x):
        features = self.backbone(x)
        features_fpn = self.neck(features)
        output = self.head(features_fpn)
        return output

modified_model = Stage2Model(model)
modified_model.eval()

再说一遍，model.eval() 很重要！

2. onnx转换：

首先指定输入：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
input1_features = torch.ones((max_pillars, 20, 6)).to(device)
input1_numPoints = torch.ones((max_pillars)).to(device)
input1_coors = torch.ones((max_pillars, 4)).to(device)

在这里是动态多输入：先构建张量，代码中的 max_pillars 是常数，然后：

with torch.no_grad():
    torch.onnx.export(model.pts_voxel_encoder,
                      (input1_features, input1_numPoints, input1_coors),
                       stage1_onnx_save_path,
                       opset_version=11,
                       do_constant_folding=True,	
                       export_params=True, 
                       verbose=False, 
                       input_names=['PFN_features', 'num_points', 'coors'],
                       output_names=['output_1_0'],                      
                       dynamic_axes={'PFN_features':{0:'effective_pillars'}, 
                                     'num_points':{0:'effective_pillars'}, 
                                     'coors':{0:'effective_pillars'}})

stage1_onnx = onnx.load(stage1_onnx_save_path)
onnx.checker.check_model(stage1_onnx)
print("Stage1 onnx model is OK!")

do_constant_folding 字段用于控制是否执行常量折叠优化；

dynamic_axes 字段用于批处理, 指定 input 和 output 的 batch（或其他维度）可变，若不想支持批处理或固定批处理大小,移除 dynamic_axes 字段即可；

check 无报错且则可视化Netron 模型结构符合你的预期，则 onnx生成本身无问题；

然后是阶段二的模型转换：

with torch.no_grad():
    torch.onnx.export(modified_model,
                    input2,
                    stage2_onnx_save_path,
                    opset_version=10,
                    do_constant_folding=True,
                    export_params=True, 
                    verbose=False, 
                    input_names=['features'],
                    output_names=['output_2_0', 'output_2_1', 'output_2_2',
                                    'output_2_3', 'output_2_4', 'output_2_5', 
                                    'output_2_6', 'output_2_7', 'output_2_8', 
                                    'output_2_9', 'output_2_10', 'output_2_11', 
                                    'output_2_12', 'output_2_13', 'output_2_14', 
                                    'output_2_15', 'output_2_16', 'output_2_17'])

model_simp2, ok2 = simplify(stage2_onnx)
assert ok2, "Simplified ONNX model could not be validated"
onnx.save(model_simp2, onnx_res_path + "stage2_test_simplifier.onnx")
print('Finished exporting simplify stage2_onnx!')

没什么好说的，输出多就多写几个；

可以选择使用 onnxsim 工具（上一篇已安装）；

3. Onnxruntime 推理及对比 pytorch 输出：

从代码的角度来讲，推理比较容易，都封装好了（以stage1为例）：

def to_numpy(exmpletensor):
    return exmpletensor.detach().cpu().numpy() if exmpletensor.requires_grad else exmpletensor.cpu().numpy()

ort_session1 = ort.InferenceSession(stage1_onnx_save_path, providers=providers)
stage1_ort_outputs = ort_session1.run(None,{'PFN_features' : to_numpy(input1_features),'num_points': to_numpy(input1_numPoints) ,'coors': to_numpy(input1_coors)})
stage1_torch_out = model.pts_voxel_encoder(input1_features, input1_numPoints, input1_coors)

得到了 onnx 侧及 pytorch 侧的输出；stage2 直接把 model 那块换成你的 modified_model；

这里留下一个疑问，希望有朋友不吝赐教：此处的 session.run 是包含数据 GPU -> CPU -> 选择推理(GPU) -> CPU? 如果想计算 纯粹GPU 推理时间该怎么弄？

对比 pytorch 输出：

np.testing.assert_allclose(to_numpy(stage1_torch_out), stage1_ort_outputs[0], rtol=1e-03, atol=1e-05)
print("Exported model1 has been tested with ONNXRuntime, and the result looks good!")

rtol 相对误差容忍度，atol 绝对误差容忍度; 如果两个元素之间的差值 <= atol + rtol * abs(stage1_ort_outputs[0])，则可接受；

4. 对比测速：

本质就是for循环多次执行推理过程，然后取平均（或者能保证GPU、CPU充分 warm-up 前提下执行一次也就差不多了）：

if test_fps:
    dummy_features, dummy_numPoints, dummy_coors = to_numpy(input1_features), to_numpy(input1_numPoints), to_numpy(input1_coors)
    
    t0 = time.perf_counter()
    for _ in range(n):
        stage1_ort_outputs = ort_session1.run(None,{'PFN_features' : dummy_features,'num_points':  dummy_numPoints,'coors': dummy_coors})
    t1 = time.perf_counter()
    print("用 onnx 完成推理1消耗的时间:%f fps:%f" % ((t1-t0)/n, n/(t1-t0)))

    t2 = time.perf_counter()
    for _ in range(n):
        stage1_torch_out = model.pts_voxel_encoder(input1_features, input1_numPoints, input1_coors)
    torch.cuda.synchronize()
    t3 = time.perf_counter()
    print("用 pytorch 完成推理1消耗的时间:%f fps:%f" % ((t3-t2)/n, n/(t3-t2)))

time.perf_counter() 非系统时钟，更加准确；

5. 结果：

pytorch: time: 0.018765, fps:53.290058；

onnruntime: time: 0.016044, fps:62.330253；

个人感觉 Onnxruntime 推理这里有一些问题，矛盾点就是上面提到的数据在GPU CPU之间的搬运；有清楚的还请私信，感激不尽！