Repository: open-mmlab/mmdetection

Issue:
NMS fails on any non-default GPU. 
I'm creating an issue that corresponds to the problem that came up in #1603. I'm reposting the body of that issue here:

I get an error when I try to run NMS code on any GPU except 0. 

The issue is that I get RuntimeError: cuda runtime error (700) : an illegal memory access was encountered at mmdet/ops/nms/src/nms_kernel.cu:103 when I try to run NMS with a Tensor on any device except CPU or 0. The error happens on this line:

  THCudaCheck(cudaMemcpy(&mask_host[0],
                        mask_dev,
                        sizeof(unsigned long long) * boxes_num * col_blocks,
                        cudaMemcpyDeviceToHost));
But I believe the issue is actually here:

  THCState *state = at::globalContext().lazyInitCUDA(); // TODO replace with getTHCState

  unsigned long long* mask_dev = NULL;
  //THCudaCheck(THCudaMalloc(state, (void**) &mask_dev,
  //                      boxes_num * col_blocks * sizeof(unsigned long long)));

  mask_dev = (unsigned long long*) THCudaMalloc(state, boxes_num * col_blocks * sizeof(unsigned long long));
My guess is that THCudaMalloc is creating the mask_dev array on device 0 and not the device corresponding to the input at::Tensor boxes. It looks like state might encode which device a new cuda array is allocated on, so my intuition would be to try and grab the state from boxes. However, I'm not a CUDA expert, so I'm probably totally off base for how to use THCState objects. I was attempting to look through the pytorch docs / source to see if I could figure something out, but I'm not having any luck.

Any pointers on how this issue might be handled would be appreciated. Note that if you have two GPUs you can reproduce the error by checking out this PR and running: xdoctest -m tests/test_nms.py test_nms_device_and_dtypes_gpu



Discussion:

Resolving patch:
diff --git a/mmdet/ops/nms/nms_wrapper.py b/mmdet/ops/nms/nms_wrapper.py
--- a/mmdet/ops/nms/nms_wrapper.py
+++ b/mmdet/ops/nms/nms_wrapper.py
@@ -21,6 +21,18 @@ def nms(dets, iou_thr, device_id=None):
     Returns:
         tuple: kept bboxes and indice, which is always the same data type as
             the input.
+
+    Example:
+        >>> dets = np.array([[49.1, 32.4, 51.0, 35.9, 0.9],
+        >>>                  [49.3, 32.9, 51.0, 35.3, 0.9],
+        >>>                  [49.2, 31.8, 51.0, 35.4, 0.5],
+        >>>                  [35.1, 11.5, 39.1, 15.7, 0.5],
+        >>>                  [35.6, 11.8, 39.3, 14.2, 0.5],
+        >>>                  [35.3, 11.5, 39.9, 14.5, 0.4],
+        >>>                  [35.2, 11.7, 39.7, 15.7, 0.3]], dtype=np.float32)
+        >>> iou_thr = 0.7
+        >>> supressed, inds = nms(dets, iou_thr)
+        >>> assert len(inds) == len(supressed) == 3
     """
     # convert dets (tensor or numpy array) to tensor
     if isinstance(dets, torch.Tensor):
@@ -50,6 +62,18 @@ def nms(dets, iou_thr, device_id=None):
 
 
 def soft_nms(dets, iou_thr, method='linear', sigma=0.5, min_score=1e-3):
+    """
+    Example:
+        >>> dets = np.array([[4., 3., 5., 3., 0.9],
+        >>>                  [4., 3., 5., 4., 0.9],
+        >>>                  [3., 1., 3., 1., 0.5],
+        >>>                  [3., 1., 3., 1., 0.5],
+        >>>                  [3., 1., 3., 1., 0.4],
+        >>>                  [3., 1., 3., 1., 0.0]], dtype=np.float32)
+        >>> iou_thr = 0.7
+        >>> supressed, inds = soft_nms(dets, iou_thr, sigma=0.5)
+        >>> assert len(inds) == len(supressed) == 3
+    """
     if isinstance(dets, torch.Tensor):
         is_tensor = True
         dets_np = dets.detach().cpu().numpy()


Test patch:
diff --git a/tests/test_nms.py b/tests/test_nms.py
new file mode 100644
--- /dev/null
+++ b/tests/test_nms.py
@@ -0,0 +1,70 @@
+"""
+CommandLine:
+    pytest tests/test_nms.py
+"""
+import numpy as np
+import torch
+
+from mmdet.ops.nms.nms_wrapper import nms
+
+
+def test_nms_device_and_dtypes_cpu():
+    """
+    CommandLine:
+        xdoctest -m tests/test_nms.py test_nms_device_and_dtypes_cpu
+    """
+    iou_thr = 0.7
+    base_dets = np.array([[49.1, 32.4, 51.0, 35.9, 0.9],
+                          [49.3, 32.9, 51.0, 35.3, 0.9],
+                          [35.3, 11.5, 39.9, 14.5, 0.4],
+                          [35.2, 11.7, 39.7, 15.7, 0.3]])
+
+    # CPU can handle float32 and float64
+    dets = base_dets.astype(np.float32)
+    supressed, inds = nms(dets, iou_thr)
+    assert dets.dtype == supressed.dtype
+    assert len(inds) == len(supressed) == 3
+
+    dets = torch.FloatTensor(base_dets)
+    surpressed, inds = nms(dets, iou_thr)
+    assert dets.dtype == surpressed.dtype
+    assert len(inds) == len(surpressed) == 3
+
+    dets = base_dets.astype(np.float64)
+    supressed, inds = nms(dets, iou_thr)
+    assert dets.dtype == supressed.dtype
+    assert len(inds) == len(supressed) == 3
+
+    dets = torch.DoubleTensor(base_dets)
+    surpressed, inds = nms(dets, iou_thr)
+    assert dets.dtype == surpressed.dtype
+    assert len(inds) == len(surpressed) == 3
+
+
+def test_nms_device_and_dtypes_gpu():
+    """
+    CommandLine:
+        xdoctest -m tests/test_nms.py test_nms_device_and_dtypes_gpu
+    """
+    if not torch.cuda.is_available():
+        import pytest
+        pytest.skip('test requires GPU and torch+cuda')
+
+    iou_thr = 0.7
+    base_dets = np.array([[49.1, 32.4, 51.0, 35.9, 0.9],
+                          [49.3, 32.9, 51.0, 35.3, 0.9],
+                          [35.3, 11.5, 39.9, 14.5, 0.4],
+                          [35.2, 11.7, 39.7, 15.7, 0.3]])
+
+    for device_id in range(torch.cuda.device_count()):
+        print('Run NMS on device_id = {!r}'.format(device_id))
+        # GPU can handle float32 but not float64
+        dets = base_dets.astype(np.float32)
+        supressed, inds = nms(dets, iou_thr, device_id)
+        assert dets.dtype == supressed.dtype
+        assert len(inds) == len(supressed) == 3
+
+        dets = torch.FloatTensor(base_dets).to(device_id)
+        surpressed, inds = nms(dets, iou_thr)
+        assert dets.dtype == surpressed.dtype
+        assert len(inds) == len(surpressed) == 3