Import pytorch model by relax frontend

Questions
#1

Hi, I’m attempting to import resnet18 from pytorch by relax.frontend, and try to verify my process by comparing the result from pytorch and relax, but got inconsist result.

The input data is from https://github.com/pytorch/vision/tree/main/gallery/. If you’re trying to run my code, you can download the assets.

The pytorch code is referenced by https://pytorch.org/vision/main/auto_examples/others/plot_scripted_tensor_transforms.html#sphx-glr-auto-examples-others-plot-scripted-tensor-transforms-py.

The relax code is referenced by https://mlc.ai/docs/get_started/tutorials/quick_start.html.

Thanks beforehand.

# import the model from pytorch

import torch
from torch import fx
from torchvision.models import resnet18, ResNet18_Weights
model = resnet18(weights=ResNet18_Weights.DEFAULT, progress=False)
model.eval()

# translate the model into relax

fx_mod = fx.symbolic_trace(model)

import tvm
from tvm import relax
from tvm.relax.frontend.torch.fx_translator import from_fx

target=tvm.target.Target("llvm")
input_info = [((1, 3, 224, 224), "float32")]
mod = from_fx(fx_mod, input_info, keep_params_as_input = True)
(mod, params) = tvm.relax.frontend.detach_params(mod)

# TVM build the model

ex = relax.build(mod, target=target)
vm = relax.VirtualMachine(ex, tvm.cpu())

# import image input

import numpy as np
import torchvision.transforms as v1
from pathlib import Path
from torchvision.io import read_image
torch.manual_seed(1)
ASSETS_PATH = Path('./assets')
dog1 = read_image(str(ASSETS_PATH / 'dog1.jpg'))

transforms = torch.nn.Sequential(
        v1.RandomCrop(224),
        v1.RandomHorizontalFlip(p=0.3),
)

device = "cpu"
weights = ResNet18_Weights.DEFAULT
model_transforms = weights.transforms(antialias=True)

dog1 = dog1.to(device)
dog1 = dog1.unsqueeze(0)
data = model_transforms(dog1)
data = data.numpy().reshape(1, 3, 224, 224).astype("float32")
data_tor = torch.Tensor(data)
data_tvm = tvm.nd.array(data)

from tvm.testing.utils import assert_allclose
assert_allclose(data_tvm.numpy(), data_tor.numpy())

# run the tvm model with the params and get the result

res_tvm = vm["main"](data_tvm, *params["main"])

# run the pytorch model

with torch.no_grad():
    res_tor = model(data_tor)

# compare the two result

res_tvm_ = np.argmax(res_tvm.numpy(), axis=1)
res_tor_ = np.argmax(res_tor.numpy(), axis=1)
print(res_tvm_, res_tor_)

import json

with open(Path('./assets') / 'imagenet_class_index.json') as labels_file:
    labels = json.load(labels_file)

print(f"TVM Prediction for Dog 1: {labels[str(res_tvm_.item())]}")
print(f"Torch Prediction for Dog 1: {labels[str(res_tor_.item())]}")

assert_allclose(res_tvm.numpy(), res_tor.numpy())
#2

It appears that my process is correct when I try to import AlexNet. However, maybe not all the operators are supported now using relax.frontend.torch.from_fx. I may test some other models soon and report the results here.

1 Like