microsoft / microsoft/onnxruntime

Will Spatial transformer operator be supported in the near future?

Open
#7,969 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

feature request
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

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 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.