xref: /aosp_15_r20/external/pytorch/scripts/onnx/test.sh (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1#!/bin/bash
2
3set -ex
4
5UNKNOWN=()
6
7# defaults
8PARALLEL=1
9export TORCH_ONNX_EXPERIMENTAL_RUNTIME_TYPE_CHECK=ERRORS
10
11while [[ $# -gt 0 ]]
12do
13    arg="$1"
14    case $arg in
15        -p|--parallel)
16            PARALLEL=1
17            shift # past argument
18            ;;
19        *) # unknown option
20            UNKNOWN+=("$1") # save it in an array for later
21            shift # past argument
22            ;;
23    esac
24done
25set -- "${UNKNOWN[@]}" # leave UNKNOWN
26
27# allows coverage to run w/o failing due to a missing plug-in
28pip install -e tools/coverage_plugins_package
29
30# realpath might not be available on MacOS
31script_path=$(python -c "import os; import sys; print(os.path.realpath(sys.argv[1]))" "${BASH_SOURCE[0]}")
32top_dir=$(dirname $(dirname $(dirname "$script_path")))
33test_paths=(
34    "$top_dir/test/onnx"
35)
36
37args=()
38args+=("-v")
39args+=("--cov")
40args+=("--cov-report")
41args+=("xml:test/coverage.xml")
42args+=("--cov-append")
43
44time python "${top_dir}/test/run_test.py" --onnx --shard "$SHARD_NUMBER" 2 --verbose
45
46if [[ "$SHARD_NUMBER" == "2" ]]; then
47  # xdoctests on onnx
48  xdoctest torch.onnx --style=google --options="+IGNORE_WHITESPACE"
49fi
50
51if [[ "$SHARD_NUMBER" == "2" ]]; then
52  # Sanity check on torchbench w/ onnx
53  pip install pandas
54  log_folder="test/.torchbench_logs"
55  device="cpu"
56  modes=("accuracy" "performance")
57  compilers=("dynamo-onnx" "torchscript-onnx")
58  suites=("huggingface" "timm_models")
59
60  mkdir -p "${log_folder}"
61  for mode in "${modes[@]}"; do
62    for compiler in "${compilers[@]}"; do
63      for suite in "${suites[@]}"; do
64        output_file="${log_folder}/${compiler}_${suite}_float32_inference_${device}_${mode}.csv"
65        bench_file="benchmarks/dynamo/${suite}.py"
66        bench_args=("--${mode}" --float32 "-d${device}" "--output=${output_file}" "--output-directory=${top_dir}" --inference -n5 "--${compiler}" --no-skip --dashboard --batch-size 1)
67        # Run only selected model for each suite to quickly validate the benchmark suite works as expected.
68        case "$suite" in
69            "torchbench")
70                bench_args+=(-k resnet18)
71                ;;
72            "huggingface")
73                bench_args+=(-k ElectraForQuestionAnswering)
74                ;;
75            "timm_models")
76                bench_args+=(-k lcnet_050)
77                ;;
78            *)
79                echo "Unknown suite: ${suite}"
80                exit 1
81                ;;
82        esac
83        python "${top_dir}/${bench_file}" "${bench_args[@]}"
84      done
85    done
86  done
87fi
88
89# Our CI expects both coverage.xml and .coverage to be within test/
90if [ -d .coverage ]; then
91  mv .coverage test/.coverage
92fi
93