Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_nms.py 2.24 KB
Newer Older
Jon Crall's avatar
Add tests for NMS / Issue with NMS on GPU 1. (#1603)  
Jon Crall committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
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