microsoft / microsoft/onnxruntime
Will Spatial transformer operator be supported in the near future?
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 8h
- Merged PRs (30d)
- 179
Description
`#!/usr/bin/env python3
-- coding: utf-8 --
import torch
import torch.nn as nn
import torch.nn.functional as F
class STNet(nn.Module):
def init(self):
super(STNet, self).init()
self.localization = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=3),
nn.MaxPool2d(2, stride=2),
nn.ReLU(True),
nn.Conv2d(32, 32, kernel_size=5),
nn.MaxPool2d(3, stride=3),
nn.ReLU(True)
)
self.fc_loc = nn.Sequential(
nn.Linear(32 * 14 * 2, 32),
nn.ReLU(True),
nn.Linear(32, 3 * 2)
)
self.fc_loc[2].weight.data.zero_()
self.fc_loc[2].bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))
def forward(self, x):
xs = self.localization(x)
xs = xs.view(-1, 32 * 14 * 2)
theta = self.fc_loc(xs)
theta = theta.view(-1, 2, 3)
x = self.f32fwd(x, theta)
return x
@torch.cuda.amp.custom_fwd(cast_inputs=torch.float32) # TODO 在 pytorch 1.6.1 中移除: https://github.com/pytorch/pytorch/issues/42218
def f32fwd(self, x, theta):
grid = F.affine_grid(theta, x.size(), align_corners=False)
x = F.grid_sample(x, grid, align_corners=False)
return x
`
The above is the Spatital transformer network, the "F.affine_grid" and "F.grid_sample" operators are not supported in the latest onnx opset version. I deeply hope STN will be supported in the next version as the importance of the transformer operator the in deep network. Besides, Can I export the STN to onnx and run inference by onnxruntime by other way in this opset version?
Thanks!
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the provided PyTorch STN example, focusing on F.affine_grid and F.grid_sample and their ONNX export behavior. Investigate the latest ONNX opset support and whether onnxruntime can execute an exported equivalent. Done would mean documented support for the operators or a verified alternative export and inference path.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 20/100