How can I apply non-maximum suppression (NMS) to a batch of images?

username: SU801T
cooked: <p>I have the following function defined for non-maximum suppression (NMS) post processing on my predictions.</p>
<p>At the moment, it is defined for a single prediction or output:</p>
<pre><code class="lang-auto">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

</code></pre>
<p>where I then apply it to a single image:</p>
<pre><code class="lang-auto">cpu_device = torch.device("cpu")

# 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)
</code></pre>
<p>However, I’m not sure how I can do this efficiently for a whole batch of images from a dataloader:</p>
<pre><code class="lang-auto">cpu_device = torch.device("cpu")
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??

</code></pre>
<p>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?</p>

username: aiGeekProgrammer
cooked: <p>Did you consider modifying <em>apply_nms</em> function so that it supports vectorized input?<br>
Like:</p>
<pre><code class="lang-auto">for prediction in orig_predictions:
   keep = torchvision.ops.nms(prediction['boxes'], ...
</code></pre>
<p>Not saying this is the best approach…</p>

username: SU801T
cooked: <p>I haven’t tried this yet but I’m assuming I would need to save the outputs for each image in a list?</p>

username: aiGeekProgrammer
cooked: <p>Well, this code from your validation loop:</p>
<pre><code class="lang-auto">  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]
</code></pre>
<p>is already processing vectorized data, i.e. many data points loaded by <em>valid_data_loader</em> into <em>images</em> and <em>targets</em> variables (depending on the <em>batch_size</em> for this loader). So doing NMS post-processing with vectorized version of <em>apply_nms</em> function should not be a problem. You have to try it to see what happens.</p>
<p>Side note: this would be not a very elegant solution, but you may set <em>batch_size</em> for <em>valid_data_loader</em> to 1 and use unchanged <em>apply_nms</em> function to see if it works at all for a single data in each validation loop and only then, maybe, move to a vectorized version. Like, you know: one small step at a time. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>

username: SU801T
cooked: <p>Thanks for the suggestions!</p>
<p>I did try it with a batch size of 1 and it does work which is useful, I’m not sure how efficient this is but it’s better than not applying NMS!</p>
<pre><code class="lang-auto"># batch size of 1 in dataloader

cpu_device = torch.device("cpu")
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]
    predictions = apply_nms(outputs[0], iou_thresh=0.3)
</code></pre>
<p>My end goal is to then find the f1-score. I don’t know if this warrants another question on the forum, but I’m not really sure how I can extract true positives, false positives etc. from the outputs…</p>