[Bug] relay.transform.AnnotateSpans does not work on modules imported from Pytorch

[gilles-reservoir]

Expected behavior

AnnotateSpans should work on modules imported from Pytorch.

Actual behavior

It fails, because the PyTorch importer creates variables with "." in their names. AnnotateSpans works by pretty-printing and then re-parsing a module. Unfortunately, the parser (tokenizer?) chokes on variables with "." in the name.

This makes it very difficult to debug other transformations when working with PyTorch models.

Environment

Ubuntu LTS, latest TVM main, torch 1.8.0

Steps to reproduce

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.dense = nn.Linear(4, 13)
        
    def forward(self, x):
        return self.dense(x)

net = Net()
data = torch.randn(17, 4)
traced = torch.jit.trace(net, data).eval()

mod, params = relay.frontend.from_pytorch(traced, [('x', tuple(data.shape))])
print(mod['main'].astext())
mod = relay.transform.AnnotateSpans()(mod)

Output:

#[version = "0.0.5"]
fn (%x: Tensor[(17, 4), float32], %dense.weight: Tensor[(13, 4), float32], %dense.bias: Tensor[(13), float32]) {
  %0 = transpose(%dense.weight, axes=[1, 0]);
  %1 = transpose(%0, axes=[1, 0]);
  %2 = nn.dense(%x, %1, units=13);
  add(%2, %dense.bias)
}
error: expected a local variable found `.`
 --> GeneratedSource:115:48
     |  
 115 |  def @main(%x: Tensor[(17, 4), float32], %dense.weight: Tensor[(13, 4), float32], %dense.bias: Tensor[(13), float32]) {
     |                                                 ^                                                                      

Unfortunately it's not totally trivial to fix this because tuples use "." for element access. Maybe the PyTorch importer / all variable creation should excise "."s from variable names, but I could see that breaking scripts that currently work.

[masahi]

Thanks I'll take a look.

[masahi]

There is a comment suggesting "x.y" should be parsed as an identifier ideally:

tvm/src/parser/parser.cc

Line 1547 in 2c67d71

// Right now we fail to parse `x.y`.

So we should fix the tokenizer / parser

[masahi]

Oh in PT models we can have a variable name like %features.1.conv.1_bias. We cannot distinguish '1' there with tuple access...

I'll write a pass to replace "." with "_" in a variable name, and let users optionally run it after PT model import.

[gilles-reservoir]

Thanks so much!