benihime91:

classes of the boxes (correct me if I am wrong)


Yes that’s the case.
I guess this will help you in what you’re looking for.

  
    
    
    Batch non-maximum suppression on the GPU vision
  
  
    hello there, 
using the awesome idea from torchvision “batched_nms”, this following code can decode for several images / several classes at once, it works because batched_nms offsets boxes according to their category, so you never perform a wrong suppression. 
I also tried to accelerate box encoding, if you are interested you can have a peek here: https://github.com/etienne87/torch_object_rnn/blob/master/core/anchors.py 
num_classes = cls_preds.shape[-1] - self.label_offset
num_anchors = box_pre…