How can I apply NMS (non-maximum suppression) on multiple images from a dataloader efficiently (PyTorch)?
I have the following function defined for non-maximum suppression (NMS) post processing on my predictions.
At the moment, it is defined for a single prediction or output:
from torchvision import transforms as torchtrans  

def apply_nms(orig_prediction, iou_thresh=0.3):
    
    # torchvision returns the indices of the bboxes to keep
    keep = torchvision.ops.nms(orig_prediction['boxes'], orig_prediction['scores'], iou_thresh)
    
    final_prediction = orig_prediction
    final_prediction['boxes'] = final_prediction['boxes'][keep]
    final_prediction['scores'] = final_prediction['scores'][keep]
    final_prediction['labels'] = final_prediction['labels'][keep]
    
    return final_prediction


where I then apply it to a single image:
cpu_device = torch.device(&quot;cpu&quot;)

# pick one image from the test set
img, target = valid_dataset[3]
# put the model in evaluation mode
model.to(cpu_device)
model.eval()
with torch.no_grad():
    output = model([img])[0]
 

nms_prediction = apply_nms(output, iou_thresh=0.1)

However, I'm not sure how I can do this efficiently for a whole batch of images from a dataloader:
cpu_device = torch.device(&quot;cpu&quot;)
model.eval()
with torch.no_grad():
  for images, targets in valid_data_loader:
    images = list(img.to(device) for img in images)
    
    outputs = model(images)
    outputs = [{k: v.to(cpu_device)for k, v in t.items()} for t in outputs]
    #DO NMS POST PROCESSING HERE??


What would be the best approach? How can I apply the above defined function for multiple images? Would this be best done in another for loop?

Have a look at the Generic Trnasform paragraph in the torchivision doc page you can use torchvision.transform.Lambda or work with functional transforms.
Here is an example with Lambda
nms_transform = torchvision.transforms.Lambda(apply_nms)

Then, you can apply the transform with the transform parameter of your dataset (or you can create your custom dataset class, as well):
dset = MyDset(..., transform=torchvision.transforms.Compose([torchvision.transforms.ToTensor(), nms_transform()])

