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