| | |
| | from copy import deepcopy |
| | from functools import partial |
| | from typing import List, Optional, Tuple |
| |
|
| | import torch |
| | import torch.nn as nn |
| | from mmdet.models.backbones.csp_darknet import Focus |
| | from mmdet.models.layers import ChannelAttention |
| | from mmengine.config import ConfigDict |
| | from torch import Tensor |
| |
|
| | from mmyolo.models import RepVGGBlock |
| | from mmyolo.models.dense_heads import (PPYOLOEHead, RTMDetHead, YOLOv5Head, |
| | YOLOv7Head, YOLOv8Head, YOLOXHead) |
| | from mmyolo.models.layers import ImplicitA, ImplicitM |
| | from ..backbone import DeployFocus, GConvFocus, NcnnFocus |
| | from ..bbox_code import (rtmdet_bbox_decoder, yolov5_bbox_decoder, |
| | yolox_bbox_decoder) |
| | from ..nms import batched_nms, efficient_nms, onnx_nms |
| | from .backend import MMYOLOBackend |
| |
|
| |
|
| | class DeployModel(nn.Module): |
| | transpose = False |
| |
|
| | def __init__(self, |
| | baseModel: nn.Module, |
| | backend: MMYOLOBackend, |
| | postprocess_cfg: Optional[ConfigDict] = None): |
| | super().__init__() |
| | self.baseModel = baseModel |
| | self.baseHead = baseModel.bbox_head |
| | self.backend = backend |
| | if postprocess_cfg is None: |
| | self.with_postprocess = False |
| | else: |
| | self.with_postprocess = True |
| | self.__init_sub_attributes() |
| | self.detector_type = type(self.baseHead) |
| | self.pre_top_k = postprocess_cfg.get('pre_top_k', 1000) |
| | self.keep_top_k = postprocess_cfg.get('keep_top_k', 100) |
| | self.iou_threshold = postprocess_cfg.get('iou_threshold', 0.65) |
| | self.score_threshold = postprocess_cfg.get('score_threshold', 0.25) |
| | self.__switch_deploy() |
| |
|
| | def __init_sub_attributes(self): |
| | self.bbox_decoder = self.baseHead.bbox_coder.decode |
| | self.prior_generate = self.baseHead.prior_generator.grid_priors |
| | self.num_base_priors = self.baseHead.num_base_priors |
| | self.featmap_strides = self.baseHead.featmap_strides |
| | self.num_classes = self.baseHead.num_classes |
| |
|
| | def __switch_deploy(self): |
| | headType = type(self.baseHead) |
| | if not self.with_postprocess: |
| | if headType in (YOLOv5Head, YOLOv7Head): |
| | self.baseHead.head_module.forward_single = self.forward_single |
| | elif headType in (PPYOLOEHead, YOLOv8Head): |
| | self.baseHead.head_module.reg_max = 0 |
| |
|
| | if self.backend in (MMYOLOBackend.HORIZONX3, MMYOLOBackend.NCNN, |
| | MMYOLOBackend.TORCHSCRIPT): |
| | self.transpose = True |
| | for layer in self.baseModel.modules(): |
| | if isinstance(layer, RepVGGBlock): |
| | layer.switch_to_deploy() |
| | elif isinstance(layer, ChannelAttention): |
| | layer.global_avgpool.forward = self.forward_gvp |
| | elif isinstance(layer, Focus): |
| | |
| | if self.backend in (MMYOLOBackend.ONNXRUNTIME, |
| | MMYOLOBackend.OPENVINO, |
| | MMYOLOBackend.TENSORRT8, |
| | MMYOLOBackend.TENSORRT7): |
| | self.baseModel.backbone.stem = DeployFocus(layer) |
| | |
| | elif self.backend == MMYOLOBackend.NCNN: |
| | self.baseModel.backbone.stem = NcnnFocus(layer) |
| | |
| | else: |
| | self.baseModel.backbone.stem = GConvFocus(layer) |
| |
|
| | def pred_by_feat(self, |
| | cls_scores: List[Tensor], |
| | bbox_preds: List[Tensor], |
| | objectnesses: Optional[List[Tensor]] = None, |
| | **kwargs): |
| | assert len(cls_scores) == len(bbox_preds) |
| | dtype = cls_scores[0].dtype |
| | device = cls_scores[0].device |
| |
|
| | nms_func = self.select_nms() |
| | if self.detector_type in (YOLOv5Head, YOLOv7Head): |
| | bbox_decoder = yolov5_bbox_decoder |
| | elif self.detector_type is RTMDetHead: |
| | bbox_decoder = rtmdet_bbox_decoder |
| | elif self.detector_type is YOLOXHead: |
| | bbox_decoder = yolox_bbox_decoder |
| | else: |
| | bbox_decoder = self.bbox_decoder |
| |
|
| | num_imgs = cls_scores[0].shape[0] |
| | featmap_sizes = [cls_score.shape[2:] for cls_score in cls_scores] |
| |
|
| | mlvl_priors = self.prior_generate( |
| | featmap_sizes, dtype=dtype, device=device) |
| |
|
| | flatten_priors = torch.cat(mlvl_priors) |
| |
|
| | mlvl_strides = [ |
| | flatten_priors.new_full( |
| | (featmap_size[0] * featmap_size[1] * self.num_base_priors, ), |
| | stride) for featmap_size, stride in zip( |
| | featmap_sizes, self.featmap_strides) |
| | ] |
| | flatten_stride = torch.cat(mlvl_strides) |
| |
|
| | |
| | flatten_cls_scores = [ |
| | cls_score.permute(0, 2, 3, 1).reshape(num_imgs, -1, |
| | self.num_classes) |
| | for cls_score in cls_scores |
| | ] |
| | cls_scores = torch.cat(flatten_cls_scores, dim=1).sigmoid() |
| |
|
| | flatten_bbox_preds = [ |
| | bbox_pred.permute(0, 2, 3, 1).reshape(num_imgs, -1, 4) |
| | for bbox_pred in bbox_preds |
| | ] |
| | flatten_bbox_preds = torch.cat(flatten_bbox_preds, dim=1) |
| |
|
| | if objectnesses is not None: |
| | flatten_objectness = [ |
| | objectness.permute(0, 2, 3, 1).reshape(num_imgs, -1) |
| | for objectness in objectnesses |
| | ] |
| | flatten_objectness = torch.cat(flatten_objectness, dim=1).sigmoid() |
| | cls_scores = cls_scores * (flatten_objectness.unsqueeze(-1)) |
| |
|
| | scores = cls_scores |
| |
|
| | bboxes = bbox_decoder(flatten_priors[None], flatten_bbox_preds, |
| | flatten_stride) |
| |
|
| | return nms_func(bboxes, scores, self.keep_top_k, self.iou_threshold, |
| | self.score_threshold, self.pre_top_k, self.keep_top_k) |
| |
|
| | def select_nms(self): |
| | if self.backend in (MMYOLOBackend.ONNXRUNTIME, MMYOLOBackend.OPENVINO): |
| | nms_func = onnx_nms |
| | elif self.backend == MMYOLOBackend.TENSORRT8: |
| | nms_func = efficient_nms |
| | elif self.backend == MMYOLOBackend.TENSORRT7: |
| | nms_func = batched_nms |
| | else: |
| | raise NotImplementedError |
| | if type(self.baseHead) in (YOLOv5Head, YOLOv7Head, YOLOXHead): |
| | nms_func = partial(nms_func, box_coding=1) |
| |
|
| | return nms_func |
| |
|
| | def forward(self, inputs: Tensor): |
| | neck_outputs = self.baseModel(inputs) |
| | if self.with_postprocess: |
| | return self.pred_by_feat(*neck_outputs) |
| | else: |
| | outputs = [] |
| | if self.transpose: |
| | for feats in zip(*neck_outputs): |
| | if self.backend in (MMYOLOBackend.NCNN, |
| | MMYOLOBackend.TORCHSCRIPT): |
| | outputs.append( |
| | torch.cat( |
| | [feat.permute(0, 2, 3, 1) for feat in feats], |
| | -1)) |
| | else: |
| | outputs.append(torch.cat(feats, 1).permute(0, 2, 3, 1)) |
| | else: |
| | for feats in zip(*neck_outputs): |
| | outputs.append(torch.cat(feats, 1)) |
| | return tuple(outputs) |
| |
|
| | @staticmethod |
| | def forward_single(x: Tensor, convs: nn.Module) -> Tuple[Tensor]: |
| | if isinstance(convs, nn.Sequential) and any( |
| | type(m) in (ImplicitA, ImplicitM) for m in convs): |
| | a, c, m = convs |
| | aw = a.implicit.clone() |
| | mw = m.implicit.clone() |
| | c = deepcopy(c) |
| | nw, cw, _, _ = c.weight.shape |
| | na, ca, _, _ = aw.shape |
| | nm, cm, _, _ = mw.shape |
| | c.bias = nn.Parameter(c.bias + ( |
| | c.weight.reshape(nw, cw) @ aw.reshape(ca, na)).squeeze(1)) |
| | c.bias = nn.Parameter(c.bias * mw.reshape(cm)) |
| | c.weight = nn.Parameter(c.weight * mw.transpose(0, 1)) |
| | convs = c |
| | feat = convs(x) |
| | return (feat, ) |
| |
|
| | @staticmethod |
| | def forward_gvp(x: Tensor) -> Tensor: |
| | return torch.mean(x, [2, 3], keepdim=True) |
| |
|
