Vectorize box encoding in FCOS

[abhi-glitchhg]

Just like #5247, i think it is possible to vectorize the following loop at this line.

  # ctrness loss
   bbox_reg_targets = [
       self.box_coder.encode_single(anchors_per_image, boxes_targets_per_image)
       for anchors_per_image, boxes_targets_per_image in zip(anchors, all_gt_boxes_targets)
   ]

Here is my attempt.

[jdsgomes]

Hi @abhi-glitchhg, thanks for working on this.
On a first pass looks good. But it is important to test that the implementation is equivalent and also test the performance, would you be able to provide these test scripts? If you need support to run benchmark in GPU for instance I am happy to help

[abhi-glitchhg]

@jdsgomes thanks for the review.

I have checked this implementation on the CPU device only.

import time
import torch
from torchvision.models.detection._utils import BoxLinearCoder

device = "cpu"
boxes = [torch.randn(3, 4).to(device=device) for i in range(5)]
rel_codes = [torch.randn( 3, 4).to(device=device) for i in range(5)]

box_coder = BoxLinearCoder(True)

# Warmup

for _ in range(1000):
    original = [
            box_coder.encode_single(bbox_regression_per_image, anchors_per_image)
            for anchors_per_image, bbox_regression_per_image in zip(boxes, rel_codes)
        ]
    original = torch.stack(original)
    new = box_coder.encode_all(rel_codes, boxes)

    torch.testing.assert_close(original, new, rtol=0, atol=1e-6)


start = time.process_time()
for _ in range(10000):
    original = [
            box_coder.encode_single(bbox_regression_per_image, anchors_per_image)
            for anchors_per_image, bbox_regression_per_image in zip(boxes, rel_codes)
        ]
    original = torch.stack(original)
print(f"time taken for encode_single: {time.process_time() -start}")

start = time.process_time()
for _ in range(10000):
    new = box_coder.encode_all(rel_codes, boxes)

print(f"time taken for encode_all: {time.process_time()-start}")

print('done ')

Running the above script, I am getting around 5x improvement.

Could you please check the same on GPU?

BTW this script is a slight modification of @datumbox's script in #6203 (review) :)

[jdsgomes]

Could you please check the same on GPU?

thanks for the confirmation. I did a benchmark using GPU and got similar results:

time taken for encode_single: 7.712065194
time taken for encode_all: 1.5909656589999983
done

[jdsgomes]

LGTM, thanks for working on this

[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