Add tests for ONNX models

[ScottTodd]

See nod-ai/SHARK-TestSuite#275

• https://github.com/onnx/models
• https://onnx.ai/models/
• https://huggingface.co/models?library=onnx

Some work has gone into importing models already at https://github.com/nod-ai/SHARK-TestSuite/tree/main/e2eshark/onnx/models

[ScottTodd]

An initial set of ONNX model tests has landed in https://github.com/iree-org/iree-test-suites/tree/main/onnx_models and is now running on every IREE PR / commit. Sample logs from a test run in IREE: https://github.com/iree-org/iree/actions/runs/11390747229/job/31693710855#step:8:19.

Test summary:

=========================== short test summary info ============================
PASSED iree-test-suites/onnx_models/tests/vision/classification_models_test.py::test_alexnet
PASSED iree-test-suites/onnx_models/tests/vision/classification_models_test.py::test_caffenet
PASSED iree-test-suites/onnx_models/tests/vision/classification_models_test.py::test_densenet_121
PASSED iree-test-suites/onnx_models/tests/vision/classification_models_test.py::test_googlenet
PASSED iree-test-suites/onnx_models/tests/vision/classification_models_test.py::test_inception_v2
PASSED iree-test-suites/onnx_models/tests/vision/classification_models_test.py::test_mnist
PASSED iree-test-suites/onnx_models/tests/vision/classification_models_test.py::test_resnet50_v1
PASSED iree-test-suites/onnx_models/tests/vision/classification_models_test.py::test_resnet50_v2
PASSED iree-test-suites/onnx_models/tests/vision/classification_models_test.py::test_shufflenet
PASSED iree-test-suites/onnx_models/tests/vision/classification_models_test.py::test_shufflenet_v2
PASSED iree-test-suites/onnx_models/tests/vision/classification_models_test.py::test_squeezenet
PASSED iree-test-suites/onnx_models/tests/vision/classification_models_test.py::test_vgg19
XFAIL iree-test-suites/onnx_models/tests/vision/classification_models_test.py::test_efficientnet_lite4
XFAIL iree-test-suites/onnx_models/tests/vision/classification_models_test.py::test_inception_v1
XFAIL iree-test-suites/onnx_models/tests/vision/classification_models_test.py::test_mobilenet
XFAIL iree-test-suites/onnx_models/tests/vision/classification_models_test.py::test_rcnn_ilsvrc13
XFAIL iree-test-suites/onnx_models/tests/vision/classification_models_test.py::test_zfnet_512
================== 12 passed, 5 xfailed in 177.35s (0:02:57) ===================

Open tasks now:

• Output and publish test report pages, like https://scotttodd.github.io/iree-test-suites/onnx_models/report_2024_09_17.html?sort=result
• Import more tests from https://github.com/onnx/models.
  • The current tests are from https://github.com/onnx/models/tree/main/validated/vision/classification.
  • See also the SHARK-TestSuite ONNX model tests at https://github.com/nod-ai/SHARK-TestSuite/tree/main/alt_e2eshark/onnx_tests/models and https://github.com/nod-ai/SHARK-TestSuite/tree/main/e2eshark/onnx/models
  • May need to edit the test util functions to support more data types. I started with just what those vision classification tests needed
• Parameterize the test suite to allow for running tests on other compiler targets and runtime HAL devices (e.g. Vulkan)

[ScottTodd]

I'll try to import more tests and parameterize across backends today.

[ScottTodd]

cc @vinayakdsci @pdhirajkumarprasad @zjgarvey

Tests are now parameterized across compiler flags (e.g. target backend choice) and runtime flags (e.g. HAL driver/device choice).

Inputs/outputs

Looking over nod-ai/SHARK-TestSuite#393 and more properties of the https://github.com/onnx/models repository, I now see that the validated model tests include reference inputs and outputs (hidden behind .tar.gz files), so we don't actually need to construct random inputs and get reference outputs from onnxruntime:

iree-test-suites/onnx_models/conftest.py

Lines 170 to 182 in 82640d6

	# Create a numpy tensor with some random data for the input.
	input_data = generate_numpy_input_for_ort_node_arg(input)
	input_data_path = onnx_path.with_name(onnx_path.stem + f"_input_{idx}.bin")
	write_ndarray_to_binary_file(input_data, input_data_path)
	inputs.append(
	IreeModelParameterMetadata(
	name=input.name,
	type=iree_type,
	data_file=input_data_path,
	)
	)
	onnx_inputs[input.name] = input_data

Input file caching

I've been thinking about caching a bit too. Right now the download code is quite naive:

iree-test-suites/onnx_models/conftest.py

Lines 278 to 284 in 82640d6

	# Download the model as needed.
	# TODO(scotttodd): move to fixture with cache / download on demand
	# TODO(scotttodd): overwrite if already existing? check SHA?
	# TODO(scotttodd): redownload if file is corrupted (e.g. partial download)
	onnx_path = test_artifacts_dir / f"{model_name}.onnx"
	if not onnx_path.exists():
	urllib.request.urlretrieve(model_url, onnx_path)

I want something similar to huggingface's cache (which is backed by git lfs), and for these files we could actually use git (lfs) directly too. Something like this:

1. Clone the onnx/models repository into a cache directory (default ~/.iree-test-suites-cache?)
2. Within that directory, run git lfs pull --include="[path to model].onnx" --exclude="" to fetch individual files, letting git handle checking if files are downloaded already
3. symlink from the repository to a test working directory

Then some extensions:

• Choose a custom dedicated cache directory with a --cache-dir flag
• Reuse an existing repository outside of a dedicated cache directory with a --onnx-models-path flag
• Cleanup cached files via https://github.com/git-lfs/git-lfs/blob/main/docs/man/git-lfs-prune.adoc (or just by deleting the whole repository or parent cache directory)

[ScottTodd]

More code to reference for caching:

• https://github.com/nod-ai/SHARK-TestSuite/blob/454722cb4213bc4cd895d152fc755a02b6d8bcf8/iree_tests/download_remote_files.py
• https://github.com/nod-ai/SHARK-TestSuite/blob/main/alt_e2eshark/onnx_tests/helper_classes.py

(But I'm liking the git repo with LFS ideas, rather than reinvent the wheel...)

[zjgarvey]

I haven't used git lfs directly much. Is there a way to pull files with git lfs on demand through a python script? It may be frustrating for someone who needs to test 10 different specific failures without downloading all models or without downloading each one individually.

[ScottTodd]

Sorta: https://stackoverflow.com/questions/74272833/how-to-git-clone-a-lfs-repo-through-python . We can try https://github.com/liberapay/git-lfs-fetch.py, but otherwise I was going to use subprocess.run().

Right now I have each test download the files it needs as the test runs. That lets you run a subset of tests and only download the files that those tests need.

In the old SHARK-TestSuite/iree_tests code I had that download_remote_files.py script that would download everything. That's useful for setting up the cache on a test runner without needing to run all of the tests or run a CI job that might otherwise hit timeouts (e.g. first run with empty cache: 1.5 hours, second run with warm cache: 10 minutes).

There are currently less than 100 tests, using less than 50GB of files. That's easy enough to manage with downloads within each test case. I would like to support all the models, like in https://github.com/nod-ai/SHARK-TestSuite/blob/main/alt_e2eshark/onnx_tests/models/external_lists/onnx_model_zoo_computer_vision_1.txt and related files, and that would pretty quickly hit some of the harder scaling challenges.