[torchlib] Migrate torchvision implementations

[justinchuby]

Adapted from https://github.com/pytorch/vision/blob/main/torchvision/ops/_register_onnx_ops.py. Also migrated tests.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 70.18%. Comparing base (6e91205) to head (30450fa).
⚠️ Report is 1 commits behind head on main.
✅ All tests successful. No failed tests found.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2569      +/-   ##
==========================================
+ Coverage   70.11%   70.18%   +0.07%     
==========================================
  Files         228      228              
  Lines       27396    27421      +25     
  Branches     2785     2786       +1     
==========================================
+ Hits        19208    19245      +37     
+ Misses       7228     7217      -11     
+ Partials      960      959       -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[Copilot]

Pull Request Overview

This PR migrates torchvision implementations by adding support for roi_align and roi_pool operations to the torch_lib vision module. The changes enable ONNX export compatibility for these computer vision operations commonly used in object detection models.

• Adds torchvision::roi_align implementation with proper coordinate transformation and sampling ratio handling
• Adds torchvision::roi_pool implementation using ONNX MaxRoiPool operator
• Updates existing torchvision::nms operation to be trace-only

[justinchuby]

@titaiwangms

[titaiwangms]

Where do we migrate this from? Was there a test we can move to here as well?

[justinchuby]

Adapted from https://github.com/pytorch/vision/blob/main/torchvision/ops/_register_onnx_ops.py. There are some tests we can move over too

[ruro]

FYI, as far as I can tell, this PR causes the following spurious warnings as long as torchvision is imported anywhere (even if the model that is being exported by the user does not contain any torchvision ops). For example:

import torch
import torchvision
torch.onnx.export(
    torch.nn.Linear(1, 1),
    torch.zeros(1, 1),
    "/tmp/linear.onnx",
)

currently leads to

W0319 05:22:31.474000 1 site-packages/torch/onnx/_internal/exporter/_schemas.py:455] Missing annotation for parameter 'input' from (input, boxes, output_size: 'Sequence[int]', spatial_scale: 'float' = 1.0, sampling_ratio: 'int' = -1, aligned: 'bool' = False). Treating as an Input.
W0319 05:22:31.475000 1 site-packages/torch/onnx/_internal/exporter/_schemas.py:455] Missing annotation for parameter 'boxes' from (input, boxes, output_size: 'Sequence[int]', spatial_scale: 'float' = 1.0, sampling_ratio: 'int' = -1, aligned: 'bool' = False). Treating as an Input.
W0319 05:22:31.475000 1 site-packages/torch/onnx/_internal/exporter/_schemas.py:455] Missing annotation for parameter 'input' from (input, boxes, output_size: 'Sequence[int]', spatial_scale: 'float' = 1.0). Treating as an Input.
W0319 05:22:31.475000 1 site-packages/torch/onnx/_internal/exporter/_schemas.py:455] Missing annotation for parameter 'boxes' from (input, boxes, output_size: 'Sequence[int]', spatial_scale: 'float' = 1.0). Treating as an Input.

[justinchuby]

FYI, as far as I can tell, this PR causes the following spurious warnings as long as torchvision is imported anywhere (even if the model that is being exported by the user does not contain any torchvision ops). For example:

import torch
import torchvision
torch.onnx.export(
    torch.nn.Linear(1, 1),
    torch.zeros(1, 1),
    "/tmp/linear.onnx",
)

currently leads to

W0319 05:22:31.474000 1 site-packages/torch/onnx/_internal/exporter/_schemas.py:455] Missing annotation for parameter 'input' from (input, boxes, output_size: 'Sequence[int]', spatial_scale: 'float' = 1.0, sampling_ratio: 'int' = -1, aligned: 'bool' = False). Treating as an Input.
W0319 05:22:31.475000 1 site-packages/torch/onnx/_internal/exporter/_schemas.py:455] Missing annotation for parameter 'boxes' from (input, boxes, output_size: 'Sequence[int]', spatial_scale: 'float' = 1.0, sampling_ratio: 'int' = -1, aligned: 'bool' = False). Treating as an Input.
W0319 05:22:31.475000 1 site-packages/torch/onnx/_internal/exporter/_schemas.py:455] Missing annotation for parameter 'input' from (input, boxes, output_size: 'Sequence[int]', spatial_scale: 'float' = 1.0). Treating as an Input.
W0319 05:22:31.475000 1 site-packages/torch/onnx/_internal/exporter/_schemas.py:455] Missing annotation for parameter 'boxes' from (input, boxes, output_size: 'Sequence[int]', spatial_scale: 'float' = 1.0). Treating as an Input.

Thanks - will find a way to hide this warning