WIP

.pre-commit-config.yaml

@@ -4,7 +4,7 @@ repos:
     hooks:
       - id: flake8
   - repo: https://github.com/timothycrosley/isort
-    rev: 5.10.1
+    rev: 4.3.21
     hooks:
       - id: isort
   - repo: https://github.com/pre-commit/mirrors-yapf

README.md

@@ -2,37 +2,33 @@
 
 ## News
 
-OpenMMLab has release [MMDeploy](https://github.com/open-mmlab/mmdeploy) which support more inference engine and repos. PRs and advices are welcome !
+This Repo now support MMDetection>=3.0
 
 ## Introduction
 
 This project aims to convert the mmdetection model to TensorRT model end2end.
 Focus on object detection for now.
 Mask support is **experiment**.
 
-support:
+Features:
 
 - fp16
 - int8(experiment)
 - batched input
 - dynamic input shape
 - combination of different modules
-- deepstream support
+- deepstream
 
 Any advices, bug reports and stars are welcome.
 
-## License
-
-This project is released under the [Apache 2.0 license](LICENSE).
-
 ## Requirement
 
-- install mmdetection:
+- install MMDetection:
 
     ```bash
     # mim is so cool!
     pip install openmim
-    mim install mmdet==2.14.0
+    mim install mmdet==3.3.0
     ```
 
 - install [torch2trt_dynamic](https://github.com/grimoire/torch2trt_dynamic):
@@ -141,10 +137,11 @@ with open(save_engine_path, mode='wb') as f:
     f.write(trt_model.state_dict()['engine'])
 ```
 
-**Note**:
-- The input of the engine is the tensor **after preprocess**.
-- The output of the engine is `num_dets, bboxes, scores, class_ids`. if you enable the `enable_mask` flag, there will be another output `mask`.
-- The bboxes output of the engine did not divided by `scale factor`.
+> \[!NOTE\]
+>
+> The input of the engine is the tensor **after preprocess**.
+> The output of the engine is `num_dets, bboxes, scores, class_ids`. if you enable the `enable_mask` flag, there will be another output `mask`.
+> The bboxes output of the engine did not divided by `scale factor`.
 
 how to use the converted model
 
@@ -218,19 +215,15 @@ Read [how-does-it-work](https://github.com/NVIDIA-AI-IOT/torch2trt#how-does-it-w
 
 Tested on:
 
-- torch=1.8.1
-- tensorrt=8.0.1.6
-- mmdetection=2.18.0
-- cuda=11.1
-
-If you find any error, please report it in the issue.
+- torch=2.2.0
+- tensorrt=8.6.1
+- mmdetection=3.3.0
+- cuda=11.7
 
 ## FAQ
 
 read [this page](./docs/FAQ.md) if you meet any problem.
 
-## Contact
-
-This repo is maintained by [@grimoire](https://github.com/grimoire)
+## License
 
-And send your resume to my e-mail if you want to join @OpenMMLab. Please read the JD for detail: [link](https://mp.weixin.qq.com/s/CzrOqITFZX-T_Kcor0hs2g)
+This project is released under the [Apache 2.0 license](LICENSE).

demo/inference.py

@@ -1,10 +1,9 @@
 from argparse import ArgumentParser
 
 import torch
-from mmdet.apis import inference_detector
-
 from mmdet2trt import mmdet2trt
 from mmdet2trt.apis import create_wrap_detector
+from mmdet.apis import inference_detector
 
 
 def main():

mmdet2trt/apis/inference.py

@@ -1,6 +1,5 @@
 import logging
 
-import mmcv
 import numpy as np
 import torch
 from addict import Addict
@@ -10,6 +9,8 @@
 from mmdet.models.roi_heads.mask_heads import FCNMaskHead
 from torch2trt_dynamic import TRTModule
 
+import mmcv
+
 logger = logging.getLogger('mmdet2trt')
 
 
@@ -183,7 +184,7 @@ def forward(self, img, img_metas, *args, **kwargs):
                 num_classes = len(self.CLASSES)
                 class_agnostic = True
                 segms_results = [[] for _ in range(num_classes)]
-                if num_dets>0:
+                if num_dets > 0:
                     for i in range(batch_size):
                         segms_results = FCNMaskHead.get_seg_masks(
                             Addict(

mmdet2trt/converters/SAConv2d.py

@@ -1,9 +1,10 @@
-import mmcv.cnn
-import mmcv.ops
 import torch
 import torch.nn.functional as F
 from torch2trt_dynamic.torch2trt_dynamic import tensorrt_converter
 
+import mmcv.cnn
+import mmcv.ops
+
 
 @tensorrt_converter('mmcv.ops.saconv.SAConv2d.forward', is_real=False)
 def convert_SAConv2d(ctx):

mmdet2trt/converters/bfp_forward.py

@@ -1,8 +1,7 @@
+import mmdet2trt.ops as mmdet2trt_ops
 import torch.nn.functional as F
 from torch2trt_dynamic.torch2trt_dynamic import tensorrt_converter
 
-import mmdet2trt.ops as mmdet2trt_ops
-
 
 @tensorrt_converter('mmdet.models.necks.BFP.forward', is_real=False)
 def convert_BFP(ctx):

mmdet2trt/converters/generalized_attention.py

@@ -1,12 +1,11 @@
 import math
 
+import mmdet2trt
 import numpy as np
 import torch
 import torch.nn.functional as F
 from torch2trt_dynamic.torch2trt_dynamic import tensorrt_converter
 
-import mmdet2trt
-
 
 def get_position_embedding(self,
                            x_q,

mmdet2trt/core/anchor/anchor_generator.py

@@ -1,7 +1,6 @@
+from mmdet2trt.models.builder import register_wrapper
 from torch import nn
 
-from mmdet2trt.models.builder import register_wraper
-
 
 class AnchorGeneratorSingle(nn.Module):
 
@@ -32,8 +31,8 @@ def forward(self, x, stride=None, device='cuda'):
             device=device)
 
 
-@register_wraper('mmdet.core.anchor.anchor_generator.YOLOAnchorGenerator')
-@register_wraper('mmdet.core.AnchorGenerator')
+@register_wrapper('mmdet.core.anchor.anchor_generator.YOLOAnchorGenerator')
+@register_wrapper('mmdet.core.AnchorGenerator')
 class AnchorGeneratorWraper(nn.Module):
 
     def __init__(self, module):
@@ -65,7 +64,7 @@ def forward(self, feat_list, device='cuda'):
         return multi_level_anchors
 
 
-@register_wraper('mmdet.core.anchor.anchor_generator.SSDAnchorGenerator')
+@register_wrapper('mmdet.core.anchor.anchor_generator.SSDAnchorGenerator')
 class SSDAnchorGeneratorWraper(nn.Module):
 
     def __init__(self, module):

mmdet2trt/core/anchor/point_generator.py

@@ -1,11 +1,10 @@
+import mmdet2trt
 import torch
+from mmdet2trt.models.builder import register_wrapper
 from torch import nn
 
-import mmdet2trt
-from mmdet2trt.models.builder import register_wraper
-
 
-@register_wraper('mmdet.core.anchor.point_generator.PointGenerator')
+@register_wrapper('mmdet.core.anchor.point_generator.PointGenerator')
 class PointGeneratorWraper(nn.Module):
 
     def __init__(self, module):
@@ -22,7 +21,7 @@ def forward(self, featmap, stride):
         return shifts
 
 
-@register_wraper('mmdet.core.anchor.point_generator.MlvlPointGenerator')
+@register_wrapper('mmdet.core.anchor.point_generator.MlvlPointGenerator')
 class MlvlPointGeneratorWraper(nn.Module):
 
     def __init__(self, module):

mmdet2trt/core/bbox/coder/bucketing_bbox_coder.py

@@ -1,10 +1,9 @@
 import numpy as np
 import torch
+from mmdet2trt.models.builder import register_wrapper
 from torch import nn
 from torch.nn import functional as F
 
-from mmdet2trt.models.builder import register_wraper
-
 from .transforms import bbox_rescale_batched
 
 
@@ -70,7 +69,7 @@ def bucket2bbox_batched(proposals,
     return bboxes, loc_confidence
 
 
-@register_wraper('mmdet.core.bbox.coder.BucketingBBoxCoder')
+@register_wrapper('mmdet.core.bbox.coder.BucketingBBoxCoder')
 class BucketingBBoxCoderWraper(nn.Module):
 
     def __init__(self, module):

mmdet2trt/core/bbox/coder/delta_xywh_bbox_coder.py

@@ -1,9 +1,8 @@
 import numpy as np
 import torch
+from mmdet2trt.models.builder import register_wrapper
 from torch import nn
 
-from mmdet2trt.models.builder import register_wraper
-
 
 def delta2bbox_custom_func(cls_scores,
                            bbox_preds,
@@ -83,7 +82,7 @@ def delta2bbox_batched(rois,
     return bboxes
 
 
-@register_wraper('mmdet.core.bbox.coder.DeltaXYWHBBoxCoder')
+@register_wrapper('mmdet.core.bbox.coder.DeltaXYWHBBoxCoder')
 class DeltaXYWHBBoxCoderWraper(nn.Module):
 
     def __init__(self, module):

mmdet2trt/core/bbox/coder/tblr_bbox_coder.py

@@ -1,8 +1,7 @@
 import torch
-from torch import nn
-
-from mmdet2trt.models.builder import register_wraper
+from mmdet2trt.models.builder import register_wrapper
 from mmdet2trt.ops import util_ops
+from torch import nn
 
 
 def batched_blr2bboxes(priors,
@@ -35,7 +34,7 @@ def batched_blr2bboxes(priors,
     return boxes
 
 
-@register_wraper('mmdet.core.bbox.coder.TBLRBBoxCoder')
+@register_wrapper('mmdet.core.bbox.coder.TBLRBBoxCoder')
 class TBLRBBoxCoderWraper(nn.Module):
 
     def __init__(self, module):

mmdet2trt/core/bbox/coder/yolo_bbox_coder.py

@@ -1,8 +1,7 @@
 import torch
+from mmdet2trt.models.builder import register_wrapper
 from torch import nn
 
-from mmdet2trt.models.builder import register_wraper
-
 
 def yolodecoder_batched(bboxes, pred_bboxes, stride):
     x_center = (bboxes[..., 0] + bboxes[..., 2]) * 0.5
@@ -23,7 +22,7 @@ def yolodecoder_batched(bboxes, pred_bboxes, stride):
     return decoded_bboxes
 
 
-@register_wraper('mmdet.core.bbox.coder.YOLOBBoxCoder')
+@register_wrapper('mmdet.core.bbox.coder.YOLOBBoxCoder')
 class YOLOBBoxCoderWraper(nn.Module):
 
     def __init__(self, module):

mmdet2trt/core/bbox/transforms.py

@@ -1,6 +1,5 @@
-import torch
-
 import mmdet2trt.ops.util_ops as mm2trt_util
+import torch
 
 
 def batched_distance2bbox(points, distance, max_shape=None):