当前位置：首页 > wzjs >正文

网站建设总体框架阿里云免费空间

wzjs 2025/9/12 8:55:17

网站建设总体框架,阿里云免费空间,小兽 wordpress,html5网站开发实战之前写过OpenPcedt下实现自定义纯点云kitti格式数据集的训练，得到了pth格式的模型。参考原文https://blog.csdn.net/m0_64293675/article/details/144294201?spm1001.2014.3001.5501 本文将给出.pth文件转换成2个onnx模型和后续onnx转plan模型的方法，便…

之前写过OpenPcedt下实现自定义纯点云kitti格式数据集的训练，得到了pth格式的模型。参考原文https://blog.csdn.net/m0_64293675/article/details/144294201?spm=1001.2014.3001.5501

本文将给出.pth文件转换成2个onnx模型和后续onnx转plan模型的方法，便于实现tensorrt的部署，【全程需要在上面文章中提到的虚拟环境(pcdet)中进行】。
参考了如下的项目：

https://github.com/CarkusL/CenterPoint/tree/main
https://github.com/NVIDIA-AI-IOT/Lidar_AI_Solution/tree/master/CUDA-CenterPoint

2个onnx模型分别是

3d稀疏卷积网络onnx模型
neck+head 网络onnx模型

首先需要新建一个文件夹，比如就叫centerpoit_export。

一、3d稀疏卷积网络onnx模型的生成

将代码拉取到本地，并将其中的det3d拷贝到centerpoit_export文件夹中

git clone https://github.com/CarkusL/CenterPoint.git
cp -r det3d/ -d centerpoit_export/

还需要修改一些代码以及编译一个扩展模块：

1、修改centerpoit_export/det3d/models/init.py
将spconv_spec = importlib.util.find_spec("spconv")注释，并改成

from importlib.util import find_spec
spconv_spec = find_spec("spconv")

在这里插入图片描述
2、修改centerpoit_export/det3d/models/backbones/scn.py
将import spconv和from spconv import SparseConv3d, SubMConv3d注释，改成

try:import spconv.pytorch as spconv from spconv.pytorch import opsfrom spconv.pytorch import SparseConv3d, SubMConv3d
except: import spconv from spconv import opsfrom spconv import SparseConv3d, SubMConv3d

在这里插入图片描述
3、编译iou3d_nms_cuda扩展模块(将3D IOU和 NMS 算子编译为 PyTorch 可调用的扩展模块)
cd 到centerpoit_export/det3d/ops/iou3d_nms路径下，将setup.py中的cuda路径修改成自己电脑中的cuda路径，然后运行下面的命令，等待编译完成。（只有编译完成了才可以import iou3d_nms_cuda）

python setup.py develop

4、拉取这个项目 https://github.com/jonygu/Lidar_AI_Solution/tree/8b71cb006d434b4c20317c66121da59b99b4508e/CUDA-CenterPoint 中的以下三个文件到centerpoit_export文件夹中
在这里插入图片描述
5、修改centerpoit_export/funcs.py

在load_scn_backbone_checkpoint函数后面新增下面的内容：

def load_scn_backbone_checkpoint_KITTI(model, file):device   = next(model.parameters()).deviceckpt     = torch.load(file, map_location=device)["model_state"]new_ckpt = collections.OrderedDict()for key, val in ckpt.items():if key.startswith("backbone_3d."):newkey = key[key.find(".")+1:]if(newkey.startswith("conv2.0.0")):newkey = "conv2.0" + newkey.split("conv2.0.0")[-1]elif(newkey.startswith("conv2.0.1")):newkey = "conv2.1" + newkey.split("conv2.0.1")[-1]elif(newkey.startswith("conv2.1")):newkey = "conv2.3" + newkey.split("conv2.1")[-1]elif(newkey.startswith("conv2.2")):newkey = "conv2.4" + newkey.split("conv2.2")[-1]elif(newkey.startswith("conv3.0.0")):newkey = "conv3.0" + newkey.split("conv3.0.0")[-1]elif(newkey.startswith("conv3.0.1")):newkey = "conv3.1" + newkey.split("conv3.0.1")[-1]elif(newkey.startswith("conv3.1")):newkey = "conv3.3" + newkey.split("conv3.1")[-1]elif(newkey.startswith("conv3.2")):newkey = "conv3.4" + newkey.split("conv3.2")[-1]elif(newkey.startswith("conv4.0.0")):newkey = "conv4.0" + newkey.split("conv4.0.0")[-1]elif(newkey.startswith("conv4.0.1")):newkey = "conv4.1" + newkey.split("conv4.0.1")[-1]elif(newkey.startswith("conv4.1")):newkey = "conv4.3" + newkey.split("conv4.1")[-1]elif(newkey.startswith("conv4.2")):newkey = "conv4.4" + newkey.split("conv4.2")[-1]elif(newkey.startswith("conv_out")):newkey = "extra_conv" + newkey.split("conv_out")[-1]else:print("backbone3d key is matching:", newkey)new_ckpt[newkey] = valmodel.load_state_dict(new_ckpt)return model

在这里插入图片描述

修改函数new_sparse_basic_block_forward中的内容

def new_sparse_basic_block_forward(self, is_fuse_relu=True):def sparse_basic_block_forward(x):identity = xout = self.conv1(x)if is_fuse_relu == False:out = out.replace_feature(self.relu(out.features))#####note train onlyout = self.conv2(out)if self.downsample is not None:identity = self.downsample(x)# if hasattr(self, 'quant_add'):#     out = out.replace_feature(self.quant_add(out.features, identity.features))# else:#     out = out.replace_feature(out.features + identity.features)  out = out.replace_feature(out.features + identity.features)            out = out.replace_feature(self.relu(out.features))return outreturn sparse_basic_block_forward

6、修改exptool.py
注释 from tools.sparseconv_quantization import QuantAdd, SparseConvolutionQunat ，以及相关的函数
在这里插入图片描述

7、修改export-scn.py
注释from tools.sparseconv_quantization import initialize, disable_quantization, quant_sparseconv_module, quant_add_module

将 main下的代码改成如下内容：

if __name__ == "__main__":parser = argparse.ArgumentParser(description="Export scn to onnx file")parser.add_argument("--in-channel", type=int, default=4, help="SCN num of input channels")parser.add_argument("--ckpt", type=str, default="(这里填.pth模型的路径,比如centerpoint.pth)", help="SCN Checkpoint (scn backbone checkpoint)")parser.add_argument("--input", type=str, default=None, help="input pickle data, random if there have no input")parser.add_argument("--save-onnx", type=str, default="(这里填转换后的.scn.onnx模型的路径+文件名，比如centerpoint_pre.csn.onnx)", help="output onnx")parser.add_argument("--save-tensor", type=str, default=None, help="Save input/output tensor to file. The purpose of this operation is to verify the inference result of c++")args = parser.parse_args()# FP16 build 稀疏模型model = SpMiddleResNetFHD(args.in_channel).cuda().eval().half()print("🔥export original model🔥") if args.ckpt:model = funcs.load_scn_backbone_checkpoint_KITTI(model, args.ckpt)# 进行层融合model = funcs.layer_fusion_bn_relu(model)         print("Fusion model:")print(model)if args.input:with open(args.input, "rb") as f:voxels, coors, spatial_shape, batch_size = pickle.load(f)voxels = torch.tensor(voxels).half().cuda()coors  = torch.tensor(coors).int().cuda()else:voxels = torch.zeros(1, args.in_channel).half().cuda()coors  = torch.zeros(1, 4).int().cuda()batch_size    = 1# spatial_shape计算公式举例：（需要根据自己训练模型时的参数去计算）# POINT_CLOUD_RANGE: [0, -50, -10, 150, 80, 10]# VOXEL_SIZE: [0.2, 0.2, 0.4]# spatial_shape = [(150-0)/0.2, (80-(-50))/0.2, (10-(-10))/0.4] = [750,650,50]     spatial_shape = [750,650,50]exptool.export_onnx(model, voxels, coors, batch_size, spatial_shape, args.save_onnx, args.save_tensor)