[3] register caffe2 mask rcnn ops in lite interpreter (#35246)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/35246
register caffe2 mask-rcnn ops in lite interpreter. It requires a leading "_" in the name.
(Note: this ignores all push blocking failures!)
Test Plan: buck build -c caffe2.expose_op_to_c10=1 //xplat/caffe2:mask_rcnn_opsAndroid
Reviewed By: iseeyuan
Differential Revision: D20528758
fbshipit-source-id: 459668a0c6cdc6aec85cb561d7acce2a5291b421
diff --git a/caffe2/operators/bbox_transform_op.cc b/caffe2/operators/bbox_transform_op.cc
index 8c17621..2a6e2d2 100644
--- a/caffe2/operators/bbox_transform_op.cc
+++ b/caffe2/operators/bbox_transform_op.cc
@@ -205,4 +205,24 @@
"Tensor output_1"
")",
BBoxTransformOpFloatCPU);
+
+ C10_EXPORT_CAFFE2_OP_TO_C10_CPU(
+ BBoxTransform2,
+ "__caffe2::BBoxTransform("
+ "Tensor rois, "
+ "Tensor deltas, "
+ "Tensor im_info, "
+ "float[] weights, "
+ "bool apply_scale, "
+ "bool rotated, "
+ "bool angle_bound_on, "
+ "int angle_bound_lo, "
+ "int angle_bound_hi, "
+ "float clip_angle_thresh, "
+ "bool legacy_plus_one"
+ ") -> ("
+ "Tensor output_0, "
+ "Tensor output_1"
+ ")",
+ BBoxTransformOpFloatCPU);
// clang-format on
diff --git a/caffe2/operators/box_with_nms_limit_op.cc b/caffe2/operators/box_with_nms_limit_op.cc
index 95b12de..f014616 100644
--- a/caffe2/operators/box_with_nms_limit_op.cc
+++ b/caffe2/operators/box_with_nms_limit_op.cc
@@ -335,4 +335,32 @@
"Tensor keeps_size"
")",
caffe2::BoxWithNMSLimitOp<caffe2::CPUContext>);
+
+C10_EXPORT_CAFFE2_OP_TO_C10_CPU(
+ BoxWithNMSLimit2,
+ "__caffe2::BoxWithNMSLimit("
+ "Tensor scores, "
+ "Tensor boxes, "
+ "Tensor batch_splits, "
+ "float score_thresh, "
+ "float nms, "
+ "int detections_per_im, "
+ "bool soft_nms_enabled, "
+ "str soft_nms_method, "
+ "float soft_nms_sigma, "
+ "float soft_nms_min_score_thres, "
+ "bool rotated, "
+ "bool cls_agnostic_bbox_reg, "
+ "bool input_boxes_include_bg_cls, "
+ "bool output_classes_include_bg_cls, "
+ "bool legacy_plus_one "
+ ") -> ("
+ "Tensor scores, "
+ "Tensor boxes, "
+ "Tensor classes, "
+ "Tensor batch_splits, "
+ "Tensor keeps, "
+ "Tensor keeps_size"
+ ")",
+ caffe2::BoxWithNMSLimitOp<caffe2::CPUContext>);
// clang-format on
diff --git a/caffe2/operators/generate_proposals_op.cc b/caffe2/operators/generate_proposals_op.cc
index 79f8560..e169a45 100644
--- a/caffe2/operators/generate_proposals_op.cc
+++ b/caffe2/operators/generate_proposals_op.cc
@@ -416,7 +416,7 @@
// clang-format off
C10_EXPORT_CAFFE2_OP_TO_C10_CPU(
- GenerateProposals,
+ GenerateProposals2,
"_caffe2::GenerateProposals("
"Tensor scores, "
"Tensor bbox_deltas, "
@@ -434,4 +434,23 @@
"bool legacy_plus_one"
") -> (Tensor output_0, Tensor output_1)",
caffe2::GenerateProposalsOp<caffe2::CPUContext>);
+C10_EXPORT_CAFFE2_OP_TO_C10_CPU(
+ GenerateProposals,
+ "__caffe2::GenerateProposals("
+ "Tensor scores, "
+ "Tensor bbox_deltas, "
+ "Tensor im_info, "
+ "Tensor anchors, "
+ "float spatial_scale, "
+ "int pre_nms_topN, "
+ "int post_nms_topN, "
+ "float nms_thresh, "
+ "float min_size, "
+ "bool angle_bound_on, "
+ "int angle_bound_lo, "
+ "int angle_bound_hi, "
+ "float clip_angle_thresh, "
+ "bool legacy_plus_one"
+ ") -> (Tensor output_0, Tensor output_1)",
+ caffe2::GenerateProposalsOp<caffe2::CPUContext>);
// clang-format on
diff --git a/caffe2/operators/heatmap_max_keypoint_op.cc b/caffe2/operators/heatmap_max_keypoint_op.cc
index 812a768..420bb89 100644
--- a/caffe2/operators/heatmap_max_keypoint_op.cc
+++ b/caffe2/operators/heatmap_max_keypoint_op.cc
@@ -171,4 +171,13 @@
"bool should_output_softmax = True"
") -> Tensor keypoints",
HeatmapMaxKeypointOpFloatCPU);
+
+C10_EXPORT_CAFFE2_OP_TO_C10_CPU(
+ HeatmapMaxKeypoint2,
+ "__caffe2::HeatmapMaxKeypoint("
+ "Tensor heatmaps, "
+ "Tensor bboxes_in, "
+ "bool should_output_softmax = True"
+ ") -> Tensor keypoints",
+ HeatmapMaxKeypointOpFloatCPU);
// clang-format on
diff --git a/caffe2/operators/roi_align_op.cc b/caffe2/operators/roi_align_op.cc
index 8d57ae4..203e771 100644
--- a/caffe2/operators/roi_align_op.cc
+++ b/caffe2/operators/roi_align_op.cc
@@ -398,4 +398,18 @@
"bool aligned"
") -> Tensor",
RoIAlignOpFloatCPU);
+
+C10_EXPORT_CAFFE2_OP_TO_C10_CPU(
+ RoIAlign2,
+ "__caffe2::RoIAlign("
+ "Tensor features, "
+ "Tensor rois, "
+ "str order, "
+ "float spatial_scale, "
+ "int pooled_h, "
+ "int pooled_w, "
+ "int sampling_ratio, "
+ "bool aligned"
+ ") -> Tensor",
+ RoIAlignOpFloatCPU);
// clang-format on