Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions configs/mmdet/detection/detection_tensorrt_static-640x640.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
_base_ = ['../_base_/base_tensorrt_static-300x300.py']

onnx_config = dict(input_shape=(640, 640))

backend_config = dict(
common_config=dict(max_workspace_size=1 << 30),
model_inputs=[
dict(
input_shapes=dict(
input=dict(
min_shape=[1, 3, 640, 640],
opt_shape=[1, 3, 640, 640],
max_shape=[1, 3, 640, 640])))
])
4 changes: 3 additions & 1 deletion mmdeploy/codebase/mmdet/models/dense_heads/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from .gfl_head import gfl_head__predict_by_feat
from .reppoints_head import reppoints_head__predict_by_feat
from .rpn_head import rpn_head__get_bboxes__ncnn, rpn_head__predict_by_feat
from .rtmdet_head import rtmdet_head__predict_by_feat
from .yolo_head import (yolov3_head__predict_by_feat,
yolov3_head__predict_by_feat__ncnn)
from .yolox_head import (yolox_head__predict_by_feat,
Expand All @@ -17,5 +18,6 @@
'yolox_head__predict_by_feat', 'base_dense_head__predict_by_feat',
'fovea_head__predict_by_feat', 'base_dense_head__predict_by_feat__ncnn',
'yolox_head__predict_by_feat__ncnn', 'gfl_head__predict_by_feat',
'reppoints_head__predict_by_feat', 'detr_head'
'reppoints_head__predict_by_feat', 'detr_head',
'rtmdet_head__predict_by_feat'
]
97 changes: 97 additions & 0 deletions mmdeploy/codebase/mmdet/models/dense_heads/rtmdet_head.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,97 @@
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Optional

import torch
from mmengine.config import ConfigDict
from mmengine.structures import InstanceData
from torch import Tensor

from mmdeploy.codebase.mmdet import get_post_processing_params
from mmdeploy.codebase.mmdet.models.layers import multiclass_nms
from mmdeploy.core import FUNCTION_REWRITER


@FUNCTION_REWRITER.register_rewriter(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since RTMDetHead share the same predict_by_feat with base_dense_head, we can export the model without this rewriter. Is there any optimization in this rewriter?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We use the same logic in YOLOX's predict_by_feat for a fair comparison with YOLO series.

func_name='mmdet.models.dense_heads.rtmdet_head.'
'RTMDetHead.predict_by_feat')
def rtmdet_head__predict_by_feat(ctx,
self,
cls_scores: List[Tensor],
bbox_preds: List[Tensor],
batch_img_metas: Optional[List[dict]] = None,
cfg: Optional[ConfigDict] = None,
rescale: bool = False,
with_nms: bool = True) -> List[InstanceData]:
"""Rewrite `predict_by_feat` of `RTMDet` for default backend.

Rewrite this function to deploy model, transform network output for a
batch into bbox predictions.

Args:
ctx: Context that contains original meta information.
cls_scores (list[Tensor]): Classification scores for all
scale levels, each is a 4D-tensor, has shape
(batch_size, num_priors * num_classes, H, W).
bbox_preds (list[Tensor]): Box energies / deltas for all
scale levels, each is a 4D-tensor, has shape
(batch_size, num_priors * 4, H, W).
batch_img_metas (list[dict], Optional): Batch image meta info.
Defaults to None.
cfg (ConfigDict, optional): Test / postprocessing
configuration, if None, test_cfg would be used.
Defaults to None.
rescale (bool): If True, return boxes in original image space.
Defaults to False.
with_nms (bool): If True, do nms before return boxes.
Defaults to True.

Returns:
tuple[Tensor, Tensor]: The first item is an (N, num_box, 5) tensor,
where 5 represent (tl_x, tl_y, br_x, br_y, score), N is batch
size and the score between 0 and 1. The shape of the second
tensor in the tuple is (N, num_box), and each element
represents the class label of the corresponding box.
"""
assert len(cls_scores) == len(bbox_preds)
device = cls_scores[0].device
cfg = self.test_cfg if cfg is None else cfg
batch_size = bbox_preds[0].shape[0]
featmap_sizes = [cls_score.shape[2:] for cls_score in cls_scores]
mlvl_priors = self.prior_generator.grid_priors(
featmap_sizes, device=device)

flatten_cls_scores = [
cls_score.permute(0, 2, 3, 1).reshape(batch_size, -1,
self.cls_out_channels)
for cls_score in cls_scores
]
flatten_bbox_preds = [
bbox_pred.permute(0, 2, 3, 1).reshape(batch_size, -1, 4)
for bbox_pred in bbox_preds
]
flatten_cls_scores = torch.cat(flatten_cls_scores, dim=1).sigmoid()
flatten_bbox_preds = torch.cat(flatten_bbox_preds, dim=1)
priors = torch.cat(mlvl_priors)
tl_x = (priors[..., 0] - flatten_bbox_preds[..., 0])
tl_y = (priors[..., 1] - flatten_bbox_preds[..., 1])
br_x = (priors[..., 0] + flatten_bbox_preds[..., 2])
br_y = (priors[..., 1] + flatten_bbox_preds[..., 3])
bboxes = torch.stack([tl_x, tl_y, br_x, br_y], -1)
# directly multiply score factor and feed to nms
max_scores, _ = torch.max(flatten_cls_scores, 1)
mask = max_scores >= cfg.score_thr
scores = flatten_cls_scores.where(mask, flatten_cls_scores.new_zeros(1))
if not with_nms:
return bboxes, scores

deploy_cfg = ctx.cfg
post_params = get_post_processing_params(deploy_cfg)
max_output_boxes_per_class = post_params.max_output_boxes_per_class
iou_threshold = cfg.nms.get('iou_threshold', post_params.iou_threshold)
score_threshold = cfg.get('score_thr', post_params.score_threshold)
pre_top_k = post_params.pre_top_k
keep_top_k = cfg.get('max_per_img', post_params.keep_top_k)

return multiclass_nms(bboxes, scores, max_output_boxes_per_class,
iou_threshold, score_threshold, pre_top_k,
keep_top_k)