[Bug][Relax][PyTorch] `grid_sample` returns the wrong shape and values after `from_exported_program`

Aperta
  • Qualcosa non va?
#20,882 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
3/5
Tempo stimato
1-2 giorni
Idoneità per principianti
62/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Attiva
Stack tecnologico
python, pytorch
Ambito
compilers, machine-learning

Direzione di ricerca

Start with the PyTorch importer's _grid_sampler_2d converter, which is reached through from_exported_program, and compare its grid argument with what relax.op.image.grid_sample expects. Run the two reproduction cases from the issue: the 1x1x4x4 input with a 1x3x5x2 grid, and the 2x2 grid. Done means the TVM output shape and values match eager PyTorch in both cases.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

needs-triage type: bug
Expected behavior

For a grid of shape [N, H, W, 2], F.grid_sample returns [N, C, H, W]. The model converted with from_exported_program should give the same shape and values as eager.

Actual behavior

TVM returns the wrong shape, with no error or warning. With a 2 x 2 grid the shape happens to be right, but the values are wrong.

eager: torch.Size([1, 1, 3, 5])
tvm:   (1, 1, 5, 2)
eager: [[[[2.5, 4.5], [10.5, 12.5]]]]
tvm:   [[[[2.5, 10.5], [12.5, 10.5]]]]

Possibly _grid_sampler_2d passes PyTorch's [N, H, W, 2] grid directly, while relax.op.image.grid_sample works on [N, 2, H, W].

Environment
  • TVM: main @ 2bf0532 (0.28.dev0), built from source, target llvm
  • PyTorch: 2.12.0+cu130
  • Python: 3.12, Ubuntu x86_64
Steps to reproduce
import torch
import tvm
import tvm_tirx
from tvm import relax
from tvm.relax.frontend.torch import from_exported_program


class M(torch.nn.Module):
    def forward(self, x, grid):
        return torch.nn.functional.grid_sample(x, grid, align_corners=False)


x = torch.arange(16, dtype=torch.float32).reshape(1, 1, 4, 4)
grid = torch.zeros(1, 3, 5, 2)
print("eager:", M()(x, grid).shape)

ep = torch.export.export(M(), (x, grid))
mod = from_exported_program(ep)
ex = tvm.compile(mod, target="llvm")
vm = relax.VirtualMachine(ex, tvm_tirx.cpu())
out = vm["main"](tvm_tirx.runtime.tensor(x.numpy()), tvm_tirx.runtime.tensor(grid.numpy()))
print("tvm:  ", out[0].numpy().shape)

# A 2 x 2 grid has the right shape by accident, but the values are wrong
grid = torch.tensor([[[[-0.5, -0.5], [0.5, -0.5]], [[-0.5, 0.5], [0.5, 0.5]]]])
print("eager:", M()(x, grid).tolist())

ep = torch.export.export(M(), (x, grid))
mod = from_exported_program(ep)
ex = tvm.compile(mod, target="llvm")
vm = relax.VirtualMachine(ex, tvm_tirx.cpu())
out = vm["main"](tvm_tirx.runtime.tensor(x.numpy()), tvm_tirx.runtime.tensor(grid.numpy()))
print("tvm:  ", out[0].numpy().tolist())
Lingua principale
Python
Stelle
13.9k
Fork
4k
Merge medio
7h 59m
PR unite (30g)
159

Preparare l'ambiente

Questo progetto non fornisce container di sviluppo, Dockerfile né guida per i contributori, quindi l'ambiente è a tuo carico: parti dal suo README e consulta la nostra guida al primo contributo per i passaggi generali.

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di apache/tvm

Tutte le issue di apache/tvm

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.