mshr-h
commented
2 weeks ago It seems like the CI is using the old docker image tlcpack/ci-cpu:20240105-165030-51bdaec6, which contains PyTorch 2.0. But the torch.export is introduced in PyTorch 2.1 so we need to upgrade it.
The dockerfile is already updated in #17338 so I guess all we have to do is to upload the new docker image to tlcpack/ci-cpu and update the image name in the jenkins file at ci/jenkins/unity_jenkinsfile.groovy.
cc @yongwww @tqchen
tqchen
Part of #17346 To keep the PR small, this implementation only supports the Core ATen operators used in AlexNet. I'll add more support after the PR is merged.
I ran the AlexNet model from Torchvision and confirmed that the inference outputs from PyTorch and Relax match. Here's a repro.
test_alexnet.py
```python import torch from torch.export import export from tvm.contrib.download import download_testdata from torchvision.models import get_model, get_model_weights from torchvision.io import read_image import tvm from tvm import relax import tvm.testing from tvm.relax.frontend.torch import from_exported_program # prepare sample image img_url = "https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true" img_name = "cat.png" img_path = download_testdata(img_url, img_name, module="data") image_tensor = read_image(img_path) # load model from torchvision model_name = "alexnet" torch_model = get_model(model_name, weights="DEFAULT").eval() weights = get_model_weights(model_name).DEFAULT transforms = weights.transforms() batch = transforms(image_tensor).unsqueeze(0) example_args = (batch,) # PyTorch exported_program = export(torch_model, args=example_args) expected: torch.Tensor = exported_program.module()(*example_args) # Relax target = "llvm" dev = tvm.cpu() mod = from_exportedprogram(exported_program) mod = tvm.relax.transform.DecomposeOpsForInference()(mod) exe = relax.build(mod, target=target) vm = relax.VirtualMachine(exe, dev) tvm_args = [tvm.nd.from_dlpack(x.contiguous()) for x in example_args] tvm_output = vm["main"](*tvm_args)[0] actual = torch.from_numpy(tvm_output.numpy()) # see if the outputs match print(f"allclose: {torch.allclose(actual, expected, rtol=1e-5, atol=1e-5, equal_nan=True)}") ```