Open vadimkantorov opened 3 years ago
Hi. Sorry for the delay.
I think this can be done with torch.vmap
. Have a look here
Have a look at pytorch/issue/42368
batched_iou = torch.vmap(box_iou)
And now we can use boxes with batches
iou = batched_iou(batched_boxes1, batched_boxes2)
I don't have much experience with vmap
but I think this would be extensible.
Can you give a run @vadimkantorov ?
This already exists in detectron2: https://github.com/facebookresearch/detectron2/blob/cbbc1ce26473cb2a5cc8f58e8ada9ae14cb41052/detectron2/structures/boxes.py#L346
But given how frequent this is, I propose this (and maybe some related box ops) are included in core torchvision
Currently it supports only
(N4, M4) -> NM
. I propose to also support(BN4, BM4) -> BNM
. It would also be nice if they supported arbitrary number of batch dimensions.This is useful for computing a cost matrix between predicted boxes and ground truth boxes for a batch of frames. Probably it can be done by adjusting tensor indexing. Something like that: