[Feature] Add carafe op for MLU (#2212)

* [Feature] Support CARAFE with Cambricon MLU backend * [Docs] Add comments for common funtions * [Test] Add allclose test for carafe * Remove print Co-authored-by: zcyKTH <zcy19950525@gmail.com> Co-authored-by: budefei <budefei@cambricon.com>
open-mmlab · Aug 19, 2022 · 832b845 · 832b845
1 parent 90e19c1
commit 832b845
Show file tree

Hide file tree

Showing 14 changed files with 1,313 additions and 5 deletions.
diff --git a/docs/en/understand_mmcv/ops.md b/docs/en/understand_mmcv/ops.md
@@ -10,7 +10,7 @@ We implement common ops used in detection, segmentation, etc.
 | BBoxOverlaps                 |     | √    | √   | √   |
 | BorderAlign                  |     | √    |     |     |
 | BoxIouRotated                | √   | √    |     |     |
-| CARAFE                       |     | √    |     |     |
+| CARAFE                       |     | √    | √   |     |
 | ChamferDistance              |     | √    |     |     |
 | CrissCrossAttention          |     | √    |     |     |
 | ContourExpand                | √   |      |     |     |

diff --git a/docs/zh_cn/understand_mmcv/ops.md b/docs/zh_cn/understand_mmcv/ops.md
@@ -10,7 +10,7 @@ MMCV 提供了检测、分割等任务中常用的算子
 | BBoxOverlaps                 |     | √    | √   | √   |
 | BorderAlign                  |     | √    |     |     |
 | BoxIouRotated                | √   | √    |     |     |
-| CARAFE                       |     | √    |     |     |
+| CARAFE                       |     | √    | √   |     |
 | ChamferDistance              |     | √    |     |     |
 | CrissCrossAttention          |     | √    |     |     |
 | ContourExpand                | √   |      |     |     |

diff --git a/mmcv/ops/carafe.py b/mmcv/ops/carafe.py
@@ -158,8 +158,6 @@ def forward(ctx, features: Tensor, masks: Tensor, kernel_size: int,
     def backward(
             ctx,
             grad_output: Tensor) -> Tuple[Tensor, Tensor, None, None, None]:
-        assert grad_output.is_cuda
-
         features, masks, rfeatures = ctx.saved_tensors
         kernel_size = ctx.kernel_size
         group_size = ctx.group_size