pytorch / pytorch/vision

Allow batch_size > 1 in quantized RoIAlign

Open
#3,655 0 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

module: ops
Dominant language
Python
Stars
17.9k
Forks
7.3k
Avg merge
1d 15h
Merged PRs (30d)
13

Description

https://github.com/pytorch/vision/pull/3624 was recently merged and we identified a potential issue: https://github.com/pytorch/vision/pull/3624#discussion_r608533402

In short, the rois tensor contains indices in the first column, but depending on the quantization, some indices cannot be properly represented. For example uneven numbers can't be represented if the tensor was quantized with qscale = 2.

To prevent any potential bug, we currently force the batch size to be 1 and hard-code the index to 0:

https://github.com/pytorch/vision/blob/07fb8ba7fad7b5b458ff862919825df4e6f60b52/torchvision/csrc/ops/quantized/cpu/qroi_align_kernel.cpp#L156-L158

https://github.com/pytorch/vision/blob/07fb8ba7fad7b5b458ff862919825df4e6f60b52/torchvision/csrc/ops/quantized/cpu/qroi_align_kernel.cpp#L39-L40

We should try to allow more than one element per batch. A potential solution would involve using per-channel quantized tensors for the roi tensor, where the first column containing the indices would be quantized in a different way from the rest of the columns.

In roi_align python op:

  • if a tensor with 5 columns is passed, raise an error if it's not per-channel: there's a high changes the indices are wrong and it's too risky. If the tensor is per-chanel, pass it through: we can assume that the user knows what they're doing and that the indices are properly represented. As a good sanity check, we can still check that the batch size is within the range of the quantized type of the first column.
  • if a list of tensors is passed, convert that list of tensors into a per-channel quantized tensor with 5 columns.

The convert_boxes_to_roi_format utils should be modified. To ensure consistency throughout the library, it should also be used in MultiScaleRoIAlign.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with torchvision/csrc/ops/quantized/cpu/qroi_align_kernel.cpp at the linked lines, then inspect the roi_align Python op, convert_boxes_to_roi_format, and MultiScaleRoIAlign. Trace how five-column ROIs and lists of tensors are handled. Done means batch sizes above one work safely with the specified per-channel checks and shared conversion behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
computer-vision, machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.