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

Monodle centerNet3D 瑞芯微RKNN、地平线Horizon芯片部署、TensorRT部署

news 2025/8/13 1:47:30

一直想做一点3D目标检测，先来一篇单目3D目标检测Monodle（基于centernet的），训练代码参考官方【代码】，这里只讲讲如何部署。

模型和完整仿真测试代码，放在github上参考链接【模型和完整代码】。

1 模型训练

训练参考官方代码 https://github.com/xinzhuma/monodle

2 导出onnx

如果按照官方代码导出的onnx，后处理写起来比较复杂且后处理时耗比较长，这里将后处理的部分代码放到模型中。原始官方导出的onnx模型如下图：
在这里插入图片描述
本示例导出的onnx模型如下图，这样导出便于写后处理代码，模型+后处理整个时耗也比较优：

把centernet3d.py 文件拷贝一份，命名为export_onnx.py，并进行如下修改：

export_onnx.py 修改后的完整代码：

import os
import cv2
import torch
import torch.nn as nn
import numpy as npfrom lib.backbones import dla
from lib.backbones.dlaup import DLAUp
from lib.backbones.hourglass import get_large_hourglass_net
from lib.backbones.hourglass import load_pretrian_modelclass CenterNet3D(nn.Module):def __init__(self, backbone='dla34', neck='DLAUp', num_class=3, downsample=4):"""CenterNet for monocular 3D object detection.:param backbone: the backbone of pipeline, such as dla34.:param neck: the necks of detection, such as dla_up.:param downsample: the ratio of down sample. [4, 8, 16, 32]:param head_conv: the channels of convolution in head. default: 256"""assert downsample in [4, 8, 16, 32]super().__init__()self.heads = {'heatmap': num_class, 'offset_2d': 2, 'size_2d': 2, 'depth': 2, 'offset_3d': 2, 'size_3d': 3,'heading': 24}self.backbone = getattr(dla, backbone)(pretrained=True, return_levels=True)channels = self.backbone.channels  # channels list for feature maps generated by backboneself.first_level = int(np.log2(downsample))scales = [2 ** i for i in range(len(channels[self.first_level:]))]self.neck = DLAUp(channels[self.first_level:], scales_list=scales)  # feature fusion [such as DLAup, FPN]self.heatmapmaxpool = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)# initialize the head of pipeline, according to heads setting.for head in self.heads.keys():if head != 'heatmap':output_channels = self.heads[head]fc = nn.Sequential(nn.Conv2d(channels[self.first_level], 256, kernel_size=3, padding=1, bias=True),nn.ReLU(inplace=True),nn.Conv2d(256, output_channels, kernel_size=1, stride=1, padding=0, bias=True))else:output_channels = self.heads[head]fc = nn.Sequential(nn.Conv2d(channels[self.first_level], 256, kernel_size=3, padding=1, bias=True),nn.ReLU(inplace=True),nn.Conv2d(256, output_channels, kernel_size=1, stride=1, padding=0, bias=True),nn.Sigmoid())# initializationif 'heatmap' in head:fc[-2].bias.data.fill_(-2.19)else:self.fill_fc_weights(fc)self.__setattr__(head, fc)def forward(self, input):feat = self.backbone(input)feat = self.neck(feat[self.first_level:])ret = {}for head in self.heads:ret[head] = self.__getattr__(head)(feat)if head == 'heatmap':heatmapmax = self.heatmapmaxpool(ret[head])ret.update({'heatmapmax': heatmapmax})return retdef fill_fc_weights(self, layers):for m in layers.modules():if isinstance(m, nn.Conv2d):nn.init.normal_(m.weight, std=0.001)if m.bias is not None:nn.init.constant_(m.bias, 0)def export_onnx(model):print('===========  onnx =========== ')dummy_input = torch.randn(1, 3, 384, 1280)input_names = ['data']output_names = ['heatmap', 'offset_2d', 'size_2d', 'depth', 'offset_3d', 'size_3d', 'heading', 'heatmapmax']torch.onnx.export(model, dummy_input, './Monodle_epoch_140.onnx', verbose=False, input_names=input_names,output_names=output_names, opset_version=11)print('======================== convert onnx Finished! .... ')if __name__ == '__main__':print('This is main ...')CLASSES = ['Pedestrian', 'Car', 'Cyclist']net = CenterNet3D(backbone='dla34')checkpoint = torch.load('./weights/checkpoint_epoch_140.pth',map_location='cpu')net.load_state_dict(checkpoint['model_state'], strict=True)net.eval()export_onnx(net)input = torch.randn((1, 3, 1280, 384))print('input1:', input.shape, input.dtype)output = net(input)print(output.keys())print(output['heatmap'].shape)print(output['offset_2d'].shape)print(output['size_2d'].shape)print(output['depth'].shape)print(output['offset_3d'].shape)print(output['size_3d'].shape)print(output['heading'].shape)