nms_rotated give wrong results on float16 tensor

[DongLiu001]

I found 'mmcv.ops.nms.nms_rotated' can't supress properly on float16 tensor. Here is my simulated result:

The first two boxes and last two boxes are overlap with iou>0.9

from mmcv.ops.nms import nms_rotated boxes = torch.tensor([ [233.2500, 276.0000, 160.0000, 50.6250, -3.0371, 0.9858], [233.6250, 276.7500, 160.8750, 52.6562, -2.9980, 0.6021], [261.0000, 326.2500, 146.6250, 46.4062, -2.9844, 0.9854], [261.0000, 326.7500, 147.3750, 47.9062, -2.9648, 0.9453] ], dtype = torch.float16, device='cuda:0') dets, keep_inds = nms_rotated(boxes[:,0:5], boxes[:,5], iou_threshold = 0.5) print(dets)

output:

tensor([[233.2500, 276.0000, 160.0000, 50.6250, -3.0371, 0.9858], [261.0000, 326.2500, 146.6250, 46.4062, -2.9844, 0.9854], [261.0000, 326.7500, 147.3750, 47.9062, -2.9648, 0.9453], [233.6250, 276.7500, 160.8750, 52.6562, -2.9980, 0.6021]], device='cuda:0', dtype=torch.float16)

If I change the dtype = torch.float, it works as expected:

output

tensor([[233.2500, 276.0000, 160.0000, 50.6250, -3.0371, 0.9858], [261.0000, 326.2500, 146.6250, 46.4062, -2.9844, 0.9854]], device='cuda:0')

It can be solved by Tensor.type(torch.float) before feed into nms_rotated. Still a bug need to be care.

[zhouzaida]

Please @zytx121 have a look.