BUG P0 (#1044)

mmdeploy/backend/coreml/ops.py

@@ -22,3 +22,66 @@ def coreml_nms(context, node):
  max_boxes=max_boxes)
 
  context.add(tuple(results), torch_name=node.outputs[0])
+
+
+@register_torch_op
+def log2(context, node):
+ """bind log2."""
+ import numpy as np
+ inputs = _get_inputs(context, node)
+ x = inputs[0]
+ log_x = mb.log(x=x)
+ context.add(mb.mul(x=log_x, y=1 / np.log(2.0)), node.name)
+
+
+@register_torch_op
+def roi_align(context, node):
+ """roi align."""
+ inputs = _get_inputs(context, node)
+
+ x = context[node.inputs[0]]
+ input_shape = x.shape # (B, C, h_in, w_in)
+ if len(input_shape) != 4:
+ raise ValueError(
+ '"CropResize" op: expected input rank 4, got {}'.format(x.rank))
+
+ const_box_info = True
+ if context[node.inputs[1]].val is None or context[
+ node.inputs[2]].val is None:
+ const_box_info = False
+
+ extrapolation_value = context[node.inputs[2]].val
+ # CoreML index information along with boxes
+ if const_box_info:
+ boxes = context[node.inputs[1]].val
+ # CoreML expects boxes/ROI in
+ # [N, 1, 5, 1, 1] format
+ boxes = boxes.reshape(boxes.shape[0], 1, boxes.shape[1], 1, 1)
+ else:
+ boxes = inputs[1]
+ boxes = mb.reshape(
+ x=boxes, shape=[boxes.shape[0], 1, boxes.shape[1], 1, 1])
+ # Get Height and Width of crop
+ h_out = inputs[3]
+ w_out = inputs[4]
+
+ # Torch input format: [B, C, h_in, w_in]
+ # CoreML input format: [B, C, h_in, w_in]
+
+ # Crop Resize
+ x = mb.crop_resize(
+ x=x,
+ roi=boxes,
+ target_height=h_out.val,
+ target_width=w_out.val,
+ normalized_coordinates=False,
+ spatial_scale=extrapolation_value,
+ box_coordinate_mode='CORNERS_WIDTH_FIRST',
+ sampling_mode='OFFSET_CORNERS',
+ )
+
+ # CoreML output format: [N, 1, C, h_out, w_out]
+ # Torch output format: [N, C, h_out, w_out]
+ x = mb.squeeze(x=x, axes=[1])
+
+ context.add(x, torch_name=node.outputs[0])

mmdeploy/codebase/mmdet/models/dense_heads/rpn_head.py

@@ -1,7 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
 
-from mmdeploy.codebase.mmdet import (get_post_processing_params,
+from mmdeploy.codebase.mmdet import (gather_topk, get_post_processing_params,
  multiclass_nms,
  pad_with_value_if_necessary)
 from mmdeploy.core import FUNCTION_REWRITER
@@ -104,11 +104,17 @@ def rpn_head__get_bboxes(ctx,
 
  if pre_topk > 0:
  _, topk_inds = scores.squeeze(2).topk(pre_topk)
- batch_inds = torch.arange(batch_size, device=device).unsqueeze(-1)
- prior_inds = topk_inds.new_zeros((1, 1))
- anchors = anchors[prior_inds, topk_inds, :]
- bbox_pred = bbox_pred[batch_inds, topk_inds, :]
- scores = scores[batch_inds, topk_inds, :]
+ bbox_pred, scores = gather_topk(
+ bbox_pred,
+ scores,
+ inds=topk_inds,
+ batch_size=batch_size,
+ is_batched=True)
+ anchors = gather_topk(
+ anchors,
+ inds=topk_inds,
+ batch_size=batch_size,
+ is_batched=False)
  mlvl_valid_bboxes.append(bbox_pred)
  mlvl_scores.append(scores)
  mlvl_valid_anchors.append(anchors)

mmdeploy/codebase/mmdet/models/roi_heads/single_level_roi_extractor.py

@@ -316,3 +316,39 @@ def single_roi_extractor__forward__openvino(ctx,
  args = (output_size, featmap_strides, sample_num, rois, *feats)
  result = SingleRoIExtractorOpenVINO.apply(*args)
  return result
+
+
+@FUNCTION_REWRITER.register_rewriter(
+ func_name='mmdet.models.roi_heads.SingleRoIExtractor.forward',
+ backend=Backend.COREML.value)
+@mark('roi_extractor', inputs=['feats', 'rois'], outputs=['bbox_feats'])
+def single_roi_extractor__forward__coreml(ctx,
+ self,
+ feats,
+ rois,
+ roi_scale_factor=None):
+ """Rewrite `forward` of SingleRoIExtractor for coreml."""
+ out_size = self.roi_layers[0].output_size
+ num_levels = len(feats)
+ roi_feats = feats[0].new_zeros(rois.shape[0], self.out_channels, *out_size)
+ if num_levels == 1:
+ assert len(rois) > 0, 'The number of rois should be positive'
+ self.roi_layers[0].use_torchvision = True
+ return self.roi_layers[0](feats[0], rois)
+
+ target_lvls = self.map_roi_levels(rois, num_levels)
+
+ if roi_scale_factor is not None:
+ rois = self.roi_rescale(rois, roi_scale_factor)
+
+ for i in range(num_levels):
+ mask = target_lvls == i
+ # inds = mask.nonzero(as_tuple=False).squeeze(1)
+ rois_t = rois * mask.unsqueeze(-1)
+ # use the roi align in torhcvision
+ self.roi_layers[i].use_torchvision = True
+ roi_feats_t = self.roi_layers[i](feats[i], rois_t)
+ roi_feats = roi_feats + roi_feats_t * (rois_t[:, -1] > 0).reshape(
+ -1, 1, 1, 1)
+ # slice to recover original size
+ return roi_feats