Change to stable sort in nms implementations

[jdsgomes]

relates to #4491
triggered by #4766

I am trying to replicate the errors and follow up on the work done in the initial PR in order to fix them and introduce the stable sort (given that is acceptably fast)

This is an investigation PR and pretty much work in progree

cc @pmeier

[facebook-github-bot]

💊 CI failures summary and remediations

As of commit 90fdb26 (more details on the Dr. CI page):

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

[jdsgomes]

Benchmark results are reassuring, the difference between using stable sort and not using seems negligible.

Code used for benchmark:

import torch
from time import time
from torchvision import models, ops


def _create_tensors_with_iou(N, iou_thresh):
    # force last box to have a pre-defined iou with the first box
    # let b0 be [x0, y0, x1, y1], and b1 be [x0, y0, x1 + d, y1],
    # then, in order to satisfy ops.iou(b0, b1) == iou_thresh,
    # we need to have d = (x1 - x0) * (1 - iou_thresh) / iou_thresh
    # Adjust the threshold upward a bit with the intent of creating
    # at least one box that exceeds (barely) the threshold and so
    # should be suppressed.
    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)
    return boxes, scores

def _create_random_boxes(N):
    boxes = torch.rand(N, 4) * 100
    scores = torch.rand(N)
    return boxes, scores

run_on="cuda"
for n in range(100, 10001, 100):
    boxes, scores = _create_random_boxes(n)
    boxes = boxes.to(device=run_on)
    scores = scores.to(device=run_on)
    start = time()
    keep = ops.nms(boxes, scores, 0.5)
    end = time()
    time_for_random = end-start
    boxes, scores = _create_tensors_with_iou(n, 0.5)
    boxes = boxes.to(device=run_on)
    scores = scores.to(device=run_on)
    start = time()
    keep = ops.nms(boxes, scores, 0.5)
    end = time()
    time_for_twiou = end - start
    print(f'{n}\t{time_for_random}\t{time_for_twiou}')

[datumbox]

Thanks @jdsgomes, this looks great to me!

This might resolve a lot of the flakiness we face on the model tests. Could you create a new issue to investigate if upon merging this, we can now remove hacks/workaround from:

vision/test/test_models.py

Lines 658 to 672 in f093d08

	# Unfortunately detection models are flaky due to the unstable sort
	# in NMS. If matching across all outputs fails, use the same approach
	# as in NMSTester.test_nms_cuda to see if this is caused by duplicate
	# scores.
	expected_file = _get_expected_file(model_name)
	expected = torch.load(expected_file)
	torch.testing.assert_close(
	output[0]["scores"], expected[0]["scores"], rtol=prec, atol=prec, check_device=False, check_dtype=False
	)
	# Note: Fmassa proposed turning off NMS by adapting the threshold
	# and then using the Hungarian algorithm as in DETR to find the
	# best match between output and expected boxes and eliminate some
	# of the flakiness. Worth exploring.
	return False # Partial validation performed

Note that it's likely that we will need to recreate the expected values of the tests. This can be part of our test improvement efforts work for next half. cc @NicolasHug @pmeier

[github-actions]

Hey @jdsgomes!

You merged this PR, but no labels were added. The list of valid labels is available at https://github.com/pytorch/vision/blob/main/.github/process_commit.py

[NicolasHug]

I'll add the BC-breaking tag just so that we remember to add a short note in the release notes. This isn't BC-breaking strictly speaking as the order of ties isn't guaranteed, but it's be nice to add a short note to let users know that they might get slightly different results in some rare cases.

Regarding the benchmark: for GPU benchmarks we should call torch.cuda.synchronize() before calling time(). This is because CUDA kernels run asynchronously, so you can reach time() before the GPU operation has actually finished.
Optionally you can also use pytorch benchmarks utils: https://pytorch.org/tutorials/recipes/recipes/benchmark.html