Batch non-maximum suppression on the GPU

username: blackberry
cooked: <p>I am trying to speed up SSD family of object detectors in PyTorch. Most implementations use a CUDA-based non-maximum suppression (NMS) for efficiency, but the implementation (from Fast/er R-CNN) is image-based. If we have batch size larger than 1, the post processing &amp; NMS becomes the bottleneck, as each image in the batch is processed serially, not in parallel, sth like this:</p>
<pre><code class="lang-auto">locs, scores = network.forward(images)         # batch of images as input
bboxes, scores = postprocess(locs, scores)   # convert bboxes to pixel coordinates

# the next part becomes the bottleneck for batched processing
for i in range(images.size(0)):   # for each image in the batch
     img_boxes = bboxes[i]
     img_scores = scores[i]
     
     for j in range(1, num_classes):
         # filter bboxes by score threshold
         keep = non_max_supression(img_boxes, img_scores, th=0.50)
         # ...
</code></pre>
<p>As the batch size increases from 1 to 4, 8,…, the post-processing &amp; NMS time per image becomes more than the <code>network.forward(images)</code>.</p>
<p>So, it is desirable to have a batched post-processing and non-maximum suppression that runs in parallel on the GPU. I could not think of a way to parallelize it in PyTorch (each image requires a separate processing), but it would be possible in CUDA.</p>
<p>I searched for such a batch implementation, but could not find any.<br>
Do you know of any batch implementation of NMS?<br>
Or any suggestions?</p>

username: Tony-Y
cooked: <p><code>batch_multiclass_non_max_suppression</code> defined in <a href="https://github.com/tensorflow/models/blob/31ae57eb8555e51493a4721f6efbea68d98a3d69/research/object_detection/core/post_processing.py" rel="nofollow noopener"> <strong>post_processing.py</strong></a> of Tensorflow Object Detection API might be helpful for the development of your NMS code.</p>

username: blackberry
cooked: <p>Thanks for the cryptic TensorFlow code link <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=6" title=":smile:" class="emoji" alt=":smile:"><br>
I will keep this in mind, but it would probably be easier to design and code it from scratch than trying to understand the cryptic TensorFlow code. I used TensorFlow for 1+ year, and I dislike it a lot.</p>

username: Tony-Y
cooked: <p><code>for i in range(images.size(0)):</code> might be parallelizable by using <a href="https://pytorch.org/docs/stable/notes/multiprocessing.html" rel="nofollow noopener">torch.multiprocessing</a>.</p>

username: blackberry
cooked: <p>Yes, but will the processes be scheduled to run in parallel on a single GPU?<br>
That is, can multiple torch processes (which are said to be drop-in replacement for Python processes) run in parallel on the same GPU? I doubt it.</p>

username: Tony-Y
cooked: <p><a href="https://pytorch.org/docs/stable/cuda.html#torch.cuda.Stream" rel="nofollow noopener">torch.cuda.Stream</a> might make it possible to execute NMS tasks on a single GPU device asynchronously.</p>
<p><a href="http://on-demand.gputechconf.com/gtc/2014/presentations/S4158-cuda-streams-best-practices-common-pitfalls.pdf" rel="nofollow noopener">A document about CUDA Stream</a></p>

username: Tony-Y
cooked: <aside class="quote" data-post="2" data-topic="4743">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alexis-jacq/48/180_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-to-use-torch-cuda-streams/4743/2">How to use torch.cuda.streams?</a> 
  </div>
  <blockquote>
    <a href="http://pytorch.org/docs/master/cuda.html?highlight=streams#torch.cuda.stream" rel="noopener nofollow ugc">torch.cuda.stream(your_stream)</a> do what you want, I guess <img width="20" height="20" src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title="wink" alt="wink" class="emoji"> 
And you can use <a href="http://pytorch.org/docs/master/notes/multiprocessing.html" rel="noopener nofollow ugc">multiprocess</a> to launch the different process with different streams. Note that multiprocess with cuda is only suported in Python 3
  </blockquote>
</aside>

<p>This comment is helpful.</p>

username: Etienne_Perot
cooked: <p>hello there,<br>
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.</p>
<p>I also tried to accelerate box encoding, if you are interested you can have a peek here: <a href="https://github.com/etienne87/torch_object_rnn/blob/master/core/anchors.py" rel="nofollow noopener">https://github.com/etienne87/torch_object_rnn/blob/master/core/anchors.py</a></p>
<pre><code class="lang-auto">num_classes = cls_preds.shape[-1] - self.label_offset
num_anchors = box_preds.shape[1]
boxes = box_preds.unsqueeze(2).expand(-1, num_anchors, num_classes, 4).contiguous()
scores = cls_preds[..., self.label_offset:].contiguous()
boxes = boxes.view(-1, 4)
scores = scores.view(-1)
rows = torch.arange(len(box_preds), dtype=torch.long)[:, None]
cols = torch.arange(num_classes, dtype=torch.long)[None, :]
idxs = rows * num_classes + cols
idxs = idxs.unsqueeze(1).expand(len(box_preds), num_anchors, num_classes)
idxs = idxs.to(scores).view(-1)
mask = scores &gt;= score_thresh
boxesf = boxes[mask].contiguous()
scoresf = scores[mask].contiguous()
idxsf = idxs[mask].contiguous()

keep = batched_nms(boxesf, scoresf, idxsf, nms_thresh)

boxes = boxesf[keep]
scores = scoresf[keep]
labels = idxsf[keep] % num_classes
batch_index = idxsf[keep] // num_classes
</code></pre>

username: elepherai
cooked: <p>Does batched_nms speed up batch inference?</p>

username: Etienne_Perot
cooked: <p>Yes, if before you were decoding sequentially (with a python for loop, image-by-image), then with batched-nms you should get a very good acceleration (for me it went from 200 + ms to 14 ms using the code above that flattens all boxes &amp; call torchvision.batched_nms)</p>

username: elepherai
cooked: <p>ok, I will try it soon, can I use in pytorch 1.1?</p>

username: Etienne_Perot
cooked: <p><img src="https://discuss.pytorch.org/images/emoji/apple/grimacing.png?v=9" title=":grimacing:" class="emoji" alt=":grimacing:"> not sure when batched_nms was introduced, i would update everything if i was you</p>

username: Etienne_Perot
cooked: <p>btw, all of it is significantly slower if you have huge number of overlapping boxes (score_threshold ~= 0.1 or less), but it should be faster than per-image decoding anyway.</p>

username: elepherai
cooked: <p>Thanks, I’m trying to change nms in maskrcnn-benchmark to batched_nms.</p>

username: Jochem_Grietens
cooked: <p>Hi All,</p>
<p>What is expected for the ‘idxs’ argument ?</p>
<p>It doesn’t seem necessary for the normal nms.</p>

username: Etienne_Perot
cooked: <p>Hello Jochem, sorry for late answer, basically batch_nms works by offsetting boxes according to the index so you don’t supress boxes that overlap while not being of same category or same image.</p>

username: bonlime
cooked: <p>Using batch NMS is actually ~30% slower for me. While the NMS itself is faster, it has to measure overlap between a larger number of boxes. Which cancels the advantage of doing it in a batched fashion</p>