Merge pull request #61 from huangyuan64/rc4main

修改knn与three_nn的适配层代码以适应新版本修改
momo609 · Oct 29, 2024 · 87f00a6 · 87f00a6
2 parents 8c45aa9 + 8f5a5b9
commit 87f00a6
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 15 deletions.
diff --git a/mmcv/ops/csrc/pytorch/npu/knn_npu.cpp b/mmcv/ops/csrc/pytorch/npu/knn_npu.cpp
@@ -12,7 +12,7 @@ void knn_forward_npu(int b, int n, int m, int nsample, const Tensor xyz,
   at::Tensor target = new_xyz.contiguous();
 
   bool is_from_knn = true;
-  EXEC_NPU_CMD(aclnnKnn, source, target, is_from_knn, dist2);
+  EXEC_NPU_CMD(aclnnKnn, source, target, is_from_knn, nsample, dist2, idx);
 }
 
 void knn_forward_impl(int b, int n, int m, int nsample, const Tensor xyz,

diff --git a/mmcv/ops/csrc/pytorch/npu/three_nn_npu.cpp b/mmcv/ops/csrc/pytorch/npu/three_nn_npu.cpp
@@ -11,7 +11,8 @@ void three_nn_forward_npu(int b, int n, int m, const Tensor unknown,
   at::Tensor target = unknown.contiguous();
 
   bool is_from_knn = false;
-  EXEC_NPU_CMD(aclnnKnn, source, target, is_from_knn, dist2);
+  int nsample = 3;
+  EXEC_NPU_CMD(aclnnKnn, source, target, is_from_knn, nsample, dist2, idx);
 }
 
 void three_nn_forward_impl(int b, int n, int m, const Tensor unknown,

diff --git a/mmcv/ops/knn.py b/mmcv/ops/knn.py
@@ -63,17 +63,10 @@ def forward(ctx,
         N = xyz.shape[1]
 
         if xyz.device.type == 'npu':
-            dist = center_xyz.new_zeros((B, npoint, N)).float()
+            dist2 = center_xyz.new_zeros((B, npoint, k)).float()
+            idx = center_xyz.new_zeros((B, npoint, k)).int()
             ext_module.knn_forward(
-                xyz,
-                center_xyz,
-                torch.Tensor([]).npu(),
-                dist,
-                b=B,
-                n=N,
-                m=npoint,
-                nsample=k)
-            dist2, idx = torch.topk(dist, k, dim=2, largest=False, sorted=True)
+                xyz, center_xyz, idx, dist2, b=B, n=N, m=npoint, nsample=k)
             zeros_idx = torch.zeros(
                 xyz.shape[0], center_xyz.shape[1], k, dtype=torch.int32).npu()
             idx.where(dist2 >= 1e10, zeros_idx)

diff --git a/mmcv/ops/three_nn.py b/mmcv/ops/three_nn.py
@@ -41,10 +41,10 @@ def forward(ctx: Any, target: torch.Tensor,
             if dtype_ == torch.float16:
                 target = target.float()
                 source = source.float()
-            dist = target.new_empty(B, N, m)
+            dist2 = target.new_empty(B, N, 3)
+            idx = target.new_empty(B, N, 3, dtype=torch.int32)
             ext_module.three_nn_forward(
-                target, source, dist, torch.Tensor([]).npu(), b=B, n=N, m=m)
-            dist2, idx = torch.topk(dist, 3, dim=2, largest=False, sorted=True)
+                target, source, dist2, idx, b=B, n=N, m=m)
             dist2 = torch.sqrt(dist2)
             if dtype_ == torch.float16:
                 dist2 = dist2.half()