Skip to content

Commit b37eeb1

Browse files
andi4191andi4191
authored
Merge pull request #750 from NVIDIA/anuragd/nox_py_test
feat: Added Python accuracy tests using Nox
2 parents 9dd4cc8 + e9865c2 commit b37eeb1

File tree

3 files changed

+178
-1
lines changed

3 files changed

+178
-1
lines changed

‎docker/setup_nox.sh

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
#!/bin/bash
2+
set -o nounset
3+
set -o errexit
4+
set -o pipefail
5+
set -e
6+
7+
post=${1:-""}
8+
9+
# fetch bazel executable
10+
BAZEL_VERSION=4.2.1
11+
ARCH=$(uname -m)
12+
if [[ "$ARCH" == "aarch64" ]]; then ARCH="arm64"; fi
13+
wget -q https://github.com/bazelbuild/bazel/releases/download/${BAZEL_VERSION}/bazel-${BAZEL_VERSION}-linux-${ARCH} -O /usr/bin/bazel
14+
chmod a+x /usr/bin/bazel
15+
export NVIDIA_TF32_OVERRIDE=0
16+
17+
cd /opt/pytorch/torch_tensorrt
18+
cp /opt/pytorch/torch_tensorrt/docker/WORKSPACE.docker /opt/pytorch/torch_tensorrt/WORKSPACE
19+
20+
pip install --user --upgrade nox
21+
TOP_DIR=/opt/pytorch/torch_tensorrt nox

‎noxfile.py

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
import nox
2+
import os
3+
4+
# Use system installed Python packages
5+
PYT_PATH='/opt/conda/lib/python3.8/site-packages' if not 'PYT_PATH' in os.environ else os.environ["PYT_PATH"]
6+
7+
# Root directory for torch_tensorrt. Set according to docker container by default
8+
TOP_DIR='/torchtrt' if not 'TOP_DIR' in os.environ else os.environ["TOP_DIR"]
9+
10+
# Download the dataset
11+
@nox.session(python=["3"], reuse_venv=True)
12+
def download_datasets(session):
13+
print("Downloading dataset to path", os.path.join(TOP_DIR, 'examples/int8/training/vgg16'))
14+
session.chdir(os.path.join(TOP_DIR, 'examples/int8/training/vgg16'))
15+
session.run_always('wget', 'https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz', external=True)
16+
session.run_always('tar', '-xvzf', 'cifar-10-binary.tar.gz', external=True)
17+
session.run_always('mkdir', '-p',
18+
os.path.join(TOP_DIR, 'tests/accuracy/datasets/data'),
19+
external=True)
20+
session.run_always('cp', '-rpf',
21+
os.path.join(TOP_DIR, 'examples/int8/training/vgg16/cifar-10-batches-bin'),
22+
os.path.join(TOP_DIR, 'tests/accuracy/datasets/data/cidar-10-batches-bin'),
23+
external=True)
24+
25+
# Download the model
26+
@nox.session(python=["3"], reuse_venv=True)
27+
def download_models(session):
28+
session.install('timm')
29+
session.chdir('tests/modules')
30+
session.run_always('python',
31+
'hub.py',
32+
env={'PYTHONPATH': PYT_PATH})
33+
34+
# Train the model
35+
@nox.session(python=["3"], reuse_venv=True)
36+
def train_model(session):
37+
session.chdir(os.path.join(TOP_DIR, 'examples/int8/training/vgg16'))
38+
session.run_always('python',
39+
'main.py',
40+
'--lr', '0.01',
41+
'--batch-size', '128',
42+
'--drop-ratio', '0.15',
43+
'--ckpt-dir', 'vgg16_ckpts',
44+
'--epochs', '25',
45+
env={'PYTHONPATH': PYT_PATH})
46+
47+
# Export model
48+
session.run_always('python',
49+
'export_ckpt.py',
50+
'vgg16_ckpts/ckpt_epoch25.pth',
51+
env={'PYTHONPATH': PYT_PATH})
52+
53+
# Finetune the model
54+
@nox.session(python=["3"], reuse_venv=True)
55+
def finetune_model(session):
56+
# Install pytorch-quantization dependency
57+
session.install('pytorch-quantization', '--extra-index-url', 'https://pypi.ngc.nvidia.com')
58+
59+
session.chdir(os.path.join(TOP_DIR, 'examples/int8/training/vgg16'))
60+
session.run_always('python',
61+
'finetune_qat.py',
62+
'--lr', '0.01',
63+
'--batch-size', '128',
64+
'--drop-ratio', '0.15',
65+
'--ckpt-dir', 'vgg16_ckpts',
66+
'--start-from', '25',
67+
'--epochs', '26',
68+
env={'PYTHONPATH': PYT_PATH})
69+
70+
# Export model
71+
session.run_always('python',
72+
'export_qat.py',
73+
'vgg16_ckpts/ckpt_epoch26.pth',
74+
env={'PYTHONPATH': PYT_PATH})
75+
76+
# Run PTQ tests
77+
@nox.session(python=["3"], reuse_venv=True)
78+
def ptq_test(session):
79+
session.chdir(os.path.join(TOP_DIR, 'tests/py'))
80+
session.run_always('cp', '-rf',
81+
os.path.join(TOP_DIR, 'examples/int8/training/vgg16', 'trained_vgg16.jit.pt'),
82+
'.',
83+
external=True)
84+
tests = [
85+
'test_ptq_dataloader_calibrator.py',
86+
'test_ptq_to_backend.py',
87+
'test_ptq_trt_calibrator.py'
88+
]
89+
for test in tests:
90+
session.run_always('python', test,
91+
env={'PYTHONPATH': PYT_PATH})
92+
93+
# Run QAT tests
94+
@nox.session(python=["3"], reuse_venv=True)
95+
def qat_test(session):
96+
session.chdir(os.path.join(TOP_DIR, 'tests/py'))
97+
session.run_always('cp', '-rf',
98+
os.path.join(TOP_DIR, 'examples/int8/training/vgg16', 'trained_vgg16_qat.jit.pt'),
99+
'.',
100+
external=True)
101+
102+
session.run_always('python',
103+
'test_qat_trt_accuracy.py',
104+
env={'PYTHONPATH': PYT_PATH})
105+
106+
# Run Python API tests
107+
@nox.session(python=["3"], reuse_venv=True)
108+
def api_test(session):
109+
session.chdir(os.path.join(TOP_DIR, 'tests/py'))
110+
tests = [
111+
"test_api.py",
112+
"test_to_backend_api.py"
113+
]
114+
for test in tests:
115+
session.run_always('python',
116+
test,
117+
env={'PYTHONPATH': PYT_PATH})
118+
119+
# Clean up
120+
@nox.session(reuse_venv=True)
121+
def cleanup(session):
122+
target = [
123+
'examples/int8/training/vgg16/*.jit.pt',
124+
'examples/int8/training/vgg16/vgg16_ckpts',
125+
'examples/int8/training/vgg16/cifar-10-*',
126+
'examples/int8/training/vgg16/data',
127+
'tests/modules/*.jit.pt',
128+
'tests/py/*.jit.pt'
129+
]
130+
131+
target = ' '.join(x for x in [os.path.join(TOP_DIR, i) for i in target])
132+
session.run_always('bash', '-c',
133+
str('rm -rf ') + target,
134+
external=True)

‎tests/README.md

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
# Tests
22

3-
Right now there are two types of tests. Converter level tests and Module level tests.
3+
Currently, following tests are supported:
4+
1. Converter level tests
5+
2. Module level tests
6+
3. Accuracy tests
47

58
The goal of Converter tests are to tests individual converters againsts specific subgraphs. The current tests in `core/conveters` are good examples on how to write these tests. In general every converter should have at least 1 test. More may be required if the operation has switches that change the behavior of the op.
69

@@ -20,6 +23,25 @@ bazel test //tests --compilation_mode=dbg --test_output=errors --jobs=4 --runs_p
2023

2124
`--jobs=4` is useful and is sometimes required to prevent too many processes to use GPU memory and cause CUDA out of memory issues.
2225

26+
Additionally, accuracy tests are supported for Python backend using Nox. Please refer [setup_nox.sh](../docker/setup_nox.sh) for reference.
27+
```
28+
# To run complete Python accuracy + API tests
29+
nox
30+
31+
# To list supported sessions
32+
nox -l
33+
```
34+
35+
Note: Supported Python tests
36+
```
37+
* download_datasets-3
38+
* download_models-3
39+
* train_model-3
40+
* finetune_model-3
41+
* ptq_test-3
42+
* qat_test-3
43+
* cleanup
44+
```
2345
### Testing using pre-built Torch-TensorRT library
2446

2547
Currently, the default strategy when we run all the tests (`bazel test //tests`) is to build the testing scripts along with the full Torch-TensorRT library (`libtorchtrt.so`) from scratch. This can lead to increased testing time and might not be needed incase you already have a pre-built Torch-TensorRT library that you want to link against.

0 commit comments

Comments
 (0)