Ask questionsknn_points unexpected error using CUDA

When using knn_points on CUDA with large K (e.g. 4000) the code fails with the error

THCudaCheck FAIL file=/opt/conda/conda-bld/pytorch_1587428111115/work/aten/src/THC/THCReduceAll.cuh line=327 error=6 : the launch timed out and was terminated
Traceback (most recent call last):
  File "", line 11, in <module>
    d = knn_points(g,p,None,None,K).dists
  File "/home/eduardo/Documents/Experiments/pytorch3d/pytorch3d/ops/", line 161, in knn_points
    p1_dists, p1_idx = _knn_points.apply(p1, p2, lengths1, lengths2, K, version)
  File "/home/eduardo/Documents/Experiments/pytorch3d/pytorch3d/ops/", line 56, in forward
    if lengths2.min() < K:
RuntimeError: cuda runtime error (6) : the launch timed out and was terminated at /opt/conda/conda-bld/pytorch_1587428111115/work/aten/src/THC/THCReduceAll.cuh:327

Tried all different versions of KNN (0 to 3) and all of them result in the same error. Using a smaller K (e.g. 2000) works just fine. If you change the device to cpu it works just fine. It does not seem to be a memory problem as the GPU memory is way below the maximum RAM of 8Gb.

Using pytorch3d compiled from source on commit 7f1e63aed1252ba8145d4a.

To Reproduce the Issue:

import torch
from pytorch3d.ops import knn_points

K = 4000
device = torch.device('cuda')
p = torch.rand(16, 200000, 3).to(device)
g = torch.rand(16, 8, 3).to(device)
d = knn_points(g, p, None, None, K).dists

Answer questions jcjohnson

Our current KNN implementation will be catastrophically slow for K=4000 -- in that regime you will probably be better off using FAISS ( for KNN instead of PyTorch3D.


Related questions

No questions were found.
Github User Rank List