@fmassa this is still WIP but a first pass would be helpful to make sure this is going in the right direction :)

As far as I can tell from a quick snippet below, the results seem to be consistent with that of the non-quantized version:

import torch
import numpy as np
from torchvision import ops

def create_tensors_with_iou(N, iou_thresh):
    # taken from test_ops. Makes sure at least one box should be removed
    # Note: this might not actually happen as we convert to ints but should be
    # the case most of the time
    boxes = torch.rand(N, 4) * 100
    boxes[:, 2:] += boxes[:, :2]
    boxes[-1, :] = boxes[0, :]
    x0, y0, x1, y1 = boxes[-1].tolist()
    iou_thresh += 1e-5
    boxes[-1, 2] += (x1 - x0) * (1 - iou_thresh) / iou_thresh
    scores = torch.rand(N) * 100
    return boxes, scores

n_removed = 0
for seed in range(1000):
    torch.manual_seed(seed)
    for iou_threshold in (.5, .8, .9):

        n_boxes = 200
        boxes, scores = create_tensors_with_iou(n_boxes, iou_threshold)

        # use integer values and clamp to the quint8 range for a fair comparison
        boxes = boxes.to(torch.int).to(torch.float).clamp(0, 255)
        scores = scores.to(torch.int).to(torch.float).clamp(0, 255)

        qboxes = torch.quantize_per_tensor(boxes, scale=1, zero_point=0, dtype=torch.quint8)
        qscores = torch.quantize_per_tensor(scores, scale=1, zero_point=0, dtype=torch.quint8)

        keep = ops.nms(boxes, scores, iou_threshold)
        qkeep = ops.nms(qboxes, qscores, iou_threshold)

        n_removed += (boxes.size(0) - keep.size(0))
        assert (keep == qkeep).all()

print(f"looks good. Removed {n_removed} boxes in total")

This is a great starting point!

A few things I think would be good to do next:

• add tests
• try seeing if it's possible to remove the dequantize_val calls, as for the operations we do here the scale and zero_point can be factored out and might get cancelled in most locations

Thanks for the review @fmassa!

Regarding using an int type higher than int8: there are different quantized dtypes, the biggest one being torch.qint32. Should we use long instead of float then?

Regarding malformed boxes and the risk of overflow, there was a similar discussion in #2582 (see in particular #2582 (comment)). I realize pytorch/pytorch#36853 was just closed so I'll try to submit another PR with the enforcement on the Python side so that the function can't be called from Python on degenerate boxes. Do we still need to take care of degenerate boxes on the C++ side?

I would recommend first opening an issue to discuss adding torch.assert_async to torchvision, as assert_async can make the CUDA context invalidated in case of asserts (so the user would need to restart the runtime).

Do we still need to take care of degenerate boxes on the C++ side?

We don't currently handle degenerate boxes at all in the implementation. My comments were about underflows that can happen due to the unsigned type being used.

Another thing to keep in mind is that if the boxes are malformed (i.e., x2 < x1) the operation will underflow and the results won't be compatible with the fp32 version anymore

@fmassa what's the rationale behind enforcing the same results between float nms and qnms on degenerate boxes?

The docstring of nms is quite clear regarding the expected format:

boxes (Tensor[N, 4])): boxes to perform NMS on. They are expected to be in (x1, y1, x2, y2) format with 0 <= x1 < x2 and 0 <= y1 < y2.

From my understanding, that means that the results are undefined on degenerate boxes, so for now I don't see a reason to enforce consistency on undefined results. But perhaps I'm missing something?

I think that the correct way to go would be to validate the boxes in Python before they're passed to the C++ functions. Whether we can use assert_async is not clear yet, but maybe we could at least do the validation for CPU tensors? That would prevent any inconsistency between float nms and qnms as qnms only supports CPUs anyway.

From my understanding, that means that the results are undefined on degenerate boxes, so for now I don't see a reason to enforce consistency on undefined results. But perhaps I'm missing something?

It's to avoid surprises to the user. But let's not worry about that now then, and just put a comment in the code saying that there might be underflow happening

I updated the PR in 7e27337 to explicitely use floats for temp variables storing the coordinates. This has a (very) minimal memory overhead over using [u]int[8] (leaving variable declarations with auto), but allows the compiler to produce vectorized instructions which significantly improves computation time. In addition this completely avoids underflows (which IMHO don't even happen anyway #3601 (comment)). On a small benchmark over 500 boxes:

nms on floats:
386 µs ± 12.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

qnms on qint32 with implicit auto declaration:
2.06 ms ± 84.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

qnms on qint32 with explicit float declration
399 µs ± 11.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Similar improvements are observed on the other quantized dtypes.

Overall I observed that qnms is slightly slower than nms, but still of the same order of magnitude. Note that it's not unexpected for a quantized version to be slower than its non-quantized counterpart. For example hardsigmoid on a 10 000 long tensor gives:

float 3.75 µs ± 50 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
quint8 5.59 µs ± 314 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
qint8 5.32 µs ± 295 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
qint32 44.9 µs ± 1.21 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)