Skip to content

Commit

Permalink
tests(//py): Restructing the nox file
Browse files Browse the repository at this point in the history 
Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>
  • Loading branch information
@narendasan
narendasan committed Jan 24, 2022
1 parent 80b906e commit 8580423
Showing 1 changed file with 238 additions and 108 deletions.
346 changes: 238 additions & 108 deletions noxfile.py
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
from distutils.command.clean import clean
import nox
import os

Expand All @@ -8,7 +9,7 @@
# TOP_DIR
TOP_DIR=os.path.dirname(os.path.realpath(__file__)) if not 'TOP_DIR' in os.environ else os.environ["TOP_DIR"]

nox.options.sessions = ["developer_tests-3"]
nox.options.sessions = ["l0_api_tests-3"]

def install_deps(session):
print("Installing deps")
Expand All @@ -30,31 +31,6 @@ def install_torch_trt(session):
session.chdir(os.path.join(TOP_DIR, "py"))
session.run("python", "setup.py", "develop")

def run_base_tests(session, use_host_env=False):
print("Running basic tests")
session.chdir(os.path.join(TOP_DIR, 'tests/py'))
tests = [
"test_api.py",
"test_to_backend_api.py"
]
for test in tests:
if use_host_env:
session.run_always('python', test, env={'PYTHONPATH': PYT_PATH})
else:
session.run_always("python", test)


# Install the latest build of torch-tensorrt
@nox.session(python=["3"], reuse_venv=True)
def developer_tests(session):
"""Basic set of tests that need to pass for code to get merged"""
install_deps(session)
install_torch_trt(session)
download_models(session)
run_base_tests(session)

# Download the dataset
@nox.session(python=["3"], reuse_venv=True)
def download_datasets(session):
print("Downloading dataset to path", os.path.join(TOP_DIR, 'examples/int8/training/vgg16'))
session.chdir(os.path.join(TOP_DIR, 'examples/int8/training/vgg16'))
Expand All @@ -68,98 +44,70 @@ def download_datasets(session):
os.path.join(TOP_DIR, 'tests/accuracy/datasets/data/cidar-10-batches-bin'),
external=True)

# Download the model
@nox.session(python=["3"], reuse_venv=True)
def download_test_models(session):
download_models(session, use_host_env=True)

# Train the model
@nox.session(python=["3"], reuse_venv=True)
def train_model(session):
def train_model(session, use_host_env=False):
session.chdir(os.path.join(TOP_DIR, 'examples/int8/training/vgg16'))
session.run_always('python',
'main.py',
'--lr', '0.01',
'--batch-size', '128',
'--drop-ratio', '0.15',
'--ckpt-dir', 'vgg16_ckpts',
'--epochs', '25',
env={'PYTHONPATH': PYT_PATH})

# Export model
session.run_always('python',
if use_host_env:
session.run_always('python',
'main.py',
'--lr', '0.01',
'--batch-size', '128',
'--drop-ratio', '0.15',
'--ckpt-dir', 'vgg16_ckpts',
'--epochs', '25',
env={'PYTHONPATH': PYT_PATH})

session.run_always('python',
'export_ckpt.py',
'vgg16_ckpts/ckpt_epoch25.pth',
env={'PYTHONPATH': PYT_PATH})
'vgg16_ckpts/ckpt_epoch25.pth')
else:
session.run_always('python',
'main.py',
'--lr', '0.01',
'--batch-size', '128',
'--drop-ratio', '0.15',
'--ckpt-dir', 'vgg16_ckpts',
'--epochs', '25')

# Finetune the model
@nox.session(python=["3"], reuse_venv=True)
def finetune_model(session):
session.run_always('python',
'export_ckpt.py',
'vgg16_ckpts/ckpt_epoch25.pth')

def finetune_model(session, use_host_env=False):
# Install pytorch-quantization dependency
session.install('pytorch-quantization', '--extra-index-url', 'https://pypi.ngc.nvidia.com')

session.chdir(os.path.join(TOP_DIR, 'examples/int8/training/vgg16'))
session.run_always('python',
'finetune_qat.py',
'--lr', '0.01',
'--batch-size', '128',
'--drop-ratio', '0.15',
'--ckpt-dir', 'vgg16_ckpts',
'--start-from', '25',
'--epochs', '26',
env={'PYTHONPATH': PYT_PATH})

# Export model
session.run_always('python',
'export_qat.py',
'vgg16_ckpts/ckpt_epoch26.pth',
env={'PYTHONPATH': PYT_PATH})

# Run PTQ tests
@nox.session(python=["3"], reuse_venv=True)
def ptq_test(session):
session.chdir(os.path.join(TOP_DIR, 'tests/py'))
session.run_always('cp', '-rf',
os.path.join(TOP_DIR, 'examples/int8/training/vgg16', 'trained_vgg16.jit.pt'),
'.',
external=True)
tests = [
'test_ptq_dataloader_calibrator.py',
'test_ptq_to_backend.py',
'test_ptq_trt_calibrator.py'
]
for test in tests:
session.run_always('python', test,
env={'PYTHONPATH': PYT_PATH})

# Run QAT tests
@nox.session(python=["3"], reuse_venv=True)
def qat_test(session):
session.chdir(os.path.join(TOP_DIR, 'tests/py'))
session.run_always('cp', '-rf',
os.path.join(TOP_DIR, 'examples/int8/training/vgg16', 'trained_vgg16_qat.jit.pt'),
'.',
external=True)

session.run_always('python',
'test_qat_trt_accuracy.py',
env={'PYTHONPATH': PYT_PATH})
if use_host_env:
session.run_always('python',
'finetune_qat.py',
'--lr', '0.01',
'--batch-size', '128',
'--drop-ratio', '0.15',
'--ckpt-dir', 'vgg16_ckpts',
'--start-from', '25',
'--epochs', '26',
env={'PYTHONPATH': PYT_PATH})

# Run Python API tests
@nox.session(python=["3"], reuse_venv=True)
def api_test(session):
session.chdir(os.path.join(TOP_DIR, 'tests/py'))
tests = [
"test_api.py",
"test_to_backend_api.py"
]
for test in tests:
# Export model
session.run_always('python',
test,
'export_qat.py',
'vgg16_ckpts/ckpt_epoch26.pth',
env={'PYTHONPATH': PYT_PATH})
else:
session.run_always('python',
'finetune_qat.py',
'--lr', '0.01',
'--batch-size', '128',
'--drop-ratio', '0.15',
'--ckpt-dir', 'vgg16_ckpts',
'--start-from', '25',
'--epochs', '26')

# Export model
session.run_always('python',
'export_qat.py',
'vgg16_ckpts/ckpt_epoch26.pth')

# Clean up
@nox.session(reuse_venv=True)
def cleanup(session):
target = [
'examples/int8/training/vgg16/*.jit.pt',
Expand All @@ -173,4 +121,186 @@ def cleanup(session):
target = ' '.join(x for x in [os.path.join(TOP_DIR, i) for i in target])
session.run_always('bash', '-c',
str('rm -rf ') + target,
external=True)
external=True)

def run_base_tests(session, use_host_env=False):
print("Running basic tests")
session.chdir(os.path.join(TOP_DIR, 'tests/py'))
tests = [
"test_api.py",
"test_to_backend_api.py"
]
for test in tests:
if use_host_env:
session.run_always('python', test, env={'PYTHONPATH': PYT_PATH})
else:
session.run_always("python", test)

def run_accuracy_tests(session, use_host_env=False):
print("Running accuracy tests")
session.chdir(os.path.join(TOP_DIR, 'tests/py'))
tests = []
for test in tests:
if use_host_env:
session.run_always('python', test, env={'PYTHONPATH': PYT_PATH})
else:
session.run_always("python", test)

def run_int8_accuracy_tests(session, use_host_env=False):
print("Running accuracy tests")
session.chdir(os.path.join(TOP_DIR, 'tests/py'))
tests = [
"test_ptq_dataloader.py",
"test_ptq_to_backend.py",
"test_qat_trt_accuracy",
]
for test in tests:
if use_host_env:
session.run_always('python', test, env={'PYTHONPATH': PYT_PATH})
else:
session.run_always("python", test)

def run_trt_compatibility_tests(session, use_host_env=False):
print("Running TensorRT compatibility tests")
session.chdir(os.path.join(TOP_DIR, 'tests/py'))
tests = [
"test_trt_intercompatibilty.py",
"test_ptq_trt_calibrator.py",
]
for test in tests:
if use_host_env:
session.run_always('python', test, env={'PYTHONPATH': PYT_PATH})
else:
session.run_always("python", test)

def run_dla_tests(session, use_host_env=False):
print("Running DLA tests")
session.chdir(os.path.join(TOP_DIR, 'tests/py'))
tests = [
"test_api_dla.py",
]
for test in tests:
if use_host_env:
session.run_always('python', test, env={'PYTHONPATH': PYT_PATH})
else:
session.run_always("python", test)

def run_multi_gpu_tests(session, use_host_env=False):
print("Running multi GPU tests")
session.chdir(os.path.join(TOP_DIR, 'tests/py'))
tests = [
"test_multi_gpu.py",
]
for test in tests:
if use_host_env:
session.run_always('python', test, env={'PYTHONPATH': PYT_PATH})
else:
session.run_always("python", test)

def run_l0_api_tests(session, use_host_env=False):
if not use_host_env:
install_deps(session)
install_torch_trt(session)
download_models(session, use_host_env)
run_base_tests(session, use_host_env)
cleanup(session)

def run_l0_dla_tests(session, use_host_env=False):
if not use_host_env:
install_deps(session)
install_torch_trt(session)
download_models(session, use_host_env)
run_base_tests(session, use_host_env)
cleanup(session)

def run_l1_accuracy_tests(session, use_host_env=False):
if not use_host_env:
install_deps(session)
install_torch_trt(session)
download_models(session, use_host_env)
download_datasets(session, use_host_env)
train_model(session, use_host_env)
run_accuracy_tests(session, use_host_env)
cleanup(session)

def run_l1_int8_accuracy_tests(session, use_host_env=False):
if not use_host_env:
install_deps(session)
install_torch_trt(session)
download_models(session, use_host_env)
download_datasets(session, use_host_env)
train_model(session, use_host_env)
finetune_model(session, use_host_env)
run_int8_accuracy_tests(session, use_host_env)
cleanup(session)

def run_l2_trt_compatibility_tests(session, use_host_env=False):
if not use_host_env:
install_deps(session)
install_torch_trt(session)
download_models(session, use_host_env)
run_trt_compatibility_tests(session, use_host_env)
cleanup(session)

def run_l2_multi_gpu_tests(session, use_host_env=False):
if not use_host_env:
install_deps(session)
install_torch_trt(session)
download_models(session, use_host_env)
run_multi_gpu_tests(session, use_host_env)
cleanup(session)

@nox.session(python=["3"], reuse_venv=True)
def l0_api_tests(session):
"""When a developer needs to check correctness for a PR or something"""
run_l0_api_tests(session, use_host_env=False)

@nox.session(python=["3"], reuse_venv=True)
def l0_api_tests_host_deps(session):
"""When a developer needs to check basic api functionality using host dependencies"""
run_l0_api_tests(session, use_host_env=True)

@nox.session(python=["3"], reuse_venv=True)
def l0_dla_tests_host_deps(session):
"""When a developer needs to check basic api functionality using host dependencies"""
run_l0_dla_tests(session, use_host_env=True)

@nox.session(python=["3"], reuse_venv=True)
def l1_accuracy_tests(session):
"""Checking accuracy performance on various usecases"""
run_l1_accuracy_tests(session, use_host_env=False)

@nox.session(python=["3"], reuse_venv=True)
def l1_accuracy_tests_host_deps(session):
"""Checking accuracy performance on various usecases using host dependencies"""
run_l1_accuracy_tests(session, use_host_env=True)

@nox.session(python=["3"], reuse_venv=True)
def l1_int8_accuracy_tests(session):
"""Checking accuracy performance on various usecases"""
run_l1_int8_accuracy_tests(session, use_host_env=False)

@nox.session(python=["3"], reuse_venv=True)
def l1_int8_accuracy_tests_host_deps(session):
"""Checking accuracy performance on various usecases using host dependencies"""
run_l1_int8_accuracy_tests(session, use_host_env=True)

@nox.session(python=["3"], reuse_venv=True)
def l2_trt_compatibility_tests(session):
"""Makes sure that TensorRT Python and Torch-TensorRT can work together"""
run_l2_trt_compatibility_tests(session, use_host_env=False)

@nox.session(python=["3"], reuse_venv=True)
def l2_trt_compatibility_tests_host_deps(session):
"""Makes sure that TensorRT Python and Torch-TensorRT can work together using host dependencies"""
run_l2_trt_compatibility_tests(session, use_host_env=True)

@nox.session(python=["3"], reuse_venv=True)
def l2_multi_gpu_tests(session):
"""Makes sure that Torch-TensorRT can operate on multi-gpu systems"""
run_l2_multi_gpu_tests(session, use_host_env=False)

@nox.session(python=["3"], reuse_venv=True)
def l2_multi_gpu_tests_host_deps(session):
"""Makes sure that Torch-TensorRT can operate on multi-gpu systems using host dependencies"""
run_l2_multi_gpu_tests(session, use_host_env=True)

0 comments on commit 8580423

Please sign in to comment.