Skip to content

Commit

Permalink
Merge pull request #700 from NVIDIA/test_global_compile
Browse files Browse the repository at this point in the history
test: Add a test for the global compile function
  • Loading branch information
narendasan authored Nov 10, 2021
2 parents 0b316d2 + 19b448c commit 8143489
Show file tree
Hide file tree
Showing 5 changed files with 118 additions and 69 deletions.
34 changes: 30 additions & 4 deletions .github/workflows/linter.yml
Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,27 @@ jobs:
cpp-linting:
name: C++ Linting
runs-on: ubuntu-latest
permissions:
actions: write
checks: write
contents: write
deployments: none
id-token: write
issues: write
discussions: write
packages: write
pull-requests: write
repository-projects: none
security-events: none
statuses: write
steps:
- uses: actions/checkout@v2
with:
ref: ${{ github.event.pull_request.head.sha }}
- name: Docker login
uses: docker/login-action@f054a8b539a109f9f41c372932f1ae047eff08c9
uses: docker/login-action@v1
with:
registry: ${{ env.REGISTRY }}
registry: ghcr.io
username: ${{ github.actor }}
password: ${{ secrets.GITHUB_TOKEN }}
- name: Run image
Expand All @@ -32,14 +45,27 @@ jobs:
py-linting:
name: Python Linting
runs-on: ubuntu-latest
permissions:
actions: write
checks: write
contents: write
deployments: none
id-token: write
issues: write
discussions: write
packages: write
pull-requests: write
repository-projects: none
security-events: none
statuses: write
steps:
- uses: actions/checkout@v2
with:
ref: ${{ github.event.pull_request.head.sha }}
- name: Docker login
uses: docker/login-action@f054a8b539a109f9f41c372932f1ae047eff08c9
uses: docker/login-action@v1
with:
registry: ${{ env.REGISTRY }}
registry: ghcr.io
username: ${{ github.actor }}
password: ${{ secrets.GITHUB_TOKEN }}
- name: Run image
Expand Down
71 changes: 45 additions & 26 deletions tests/py/test_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,9 +31,25 @@ def test_compile_traced(self):

def test_compile_script(self):
trt_mod = torchtrt.ts.compile(self.scripted_model,
inputs=[self.input],
device=torchtrt.Device(gpu_id=0),
enabled_precisions={torch.float})
inputs=[self.input],
device=torchtrt.Device(gpu_id=0),
enabled_precisions={torch.float})
same = (trt_mod(self.input) - self.scripted_model(self.input)).abs().max()
self.assertTrue(same < 2e-2)

def test_compile_global(self):
trt_mod = torchtrt.compile(self.scripted_model,
inputs=[self.input],
device=torchtrt.Device(gpu_id=0),
enabled_precisions={torch.float})
same = (trt_mod(self.input) - self.scripted_model(self.input)).abs().max()
self.assertTrue(same < 2e-2)

def test_compile_global_nn_mod(self):
trt_mod = torchtrt.compile(self.model,
inputs=[self.input],
device=torchtrt.Device(gpu_id=0),
enabled_precisions={torch.float})
same = (trt_mod(self.input) - self.scripted_model(self.input)).abs().max()
self.assertTrue(same < 2e-2)

Expand Down Expand Up @@ -207,16 +223,16 @@ def setUp(self):
def test_input_use_default_fp32(self):
ts_model = torch.jit.script(self.model)
trt_mod = torchtrt.ts.compile(ts_model,
inputs=[torchtrt.Input(self.input.shape)],
enabled_precisions={torch.float, torch.half})
inputs=[torchtrt.Input(self.input.shape)],
enabled_precisions={torch.float, torch.half})
trt_mod(self.input)

def test_input_respect_user_setting_fp32_weights_fp16_in(self):
ts_model = torch.jit.script(self.model)
trt_mod = torchtrt.ts.compile(ts_model,
inputs=[self.input.half()],
require_full_compilation=True,
enabled_precisions={torch.float, torch.half})
inputs=[self.input.half()],
require_full_compilation=True,
enabled_precisions={torch.float, torch.half})
trt_mod(self.input.half())

def test_input_respect_user_setting_fp32_weights_fp16_in_non_constructor(self):
Expand All @@ -225,9 +241,9 @@ def test_input_respect_user_setting_fp32_weights_fp16_in_non_constructor(self):
input_spec.dtype = torch.half

trt_mod = torchtrt.ts.compile(ts_model,
inputs=[input_spec],
require_full_compilation=True,
enabled_precisions={torch.float, torch.half})
inputs=[input_spec],
require_full_compilation=True,
enabled_precisions={torch.float, torch.half})
trt_mod(self.input.half())


Expand All @@ -241,8 +257,8 @@ def test_input_use_default_fp16(self):
half_mod.half()

trt_mod = torchtrt.ts.compile(half_mod,
inputs=[torchtrt.Input(self.input.shape)],
enabled_precisions={torch.float, torch.half})
inputs=[torchtrt.Input(self.input.shape)],
enabled_precisions={torch.float, torch.half})
trt_mod(self.input.half())

def test_input_use_default_fp16_without_fp16_enabled(self):
Expand All @@ -257,9 +273,9 @@ def test_input_respect_user_setting_fp16_weights_fp32_in(self):
half_mod.half()

trt_mod = torchtrt.ts.compile(half_mod,
inputs=[self.input],
require_full_compilation=True,
enabled_precisions={torch.float, torch.half})
inputs=[self.input],
require_full_compilation=True,
enabled_precisions={torch.float, torch.half})
trt_mod(self.input)

def test_input_respect_user_setting_fp16_weights_fp32_in_non_constuctor(self):
Expand All @@ -270,9 +286,9 @@ def test_input_respect_user_setting_fp16_weights_fp32_in_non_constuctor(self):
input_spec.dtype = torch.float

trt_mod = torchtrt.ts.compile(half_mod,
inputs=[input_spec],
require_full_compilation=True,
enabled_precisions={torch.float, torch.half})
inputs=[input_spec],
require_full_compilation=True,
enabled_precisions={torch.float, torch.half})
trt_mod(self.input)


Expand Down Expand Up @@ -352,14 +368,15 @@ def test_from_torch(self):
self.assertEqual(device.device_type, torchtrt.DeviceType.GPU)
self.assertEqual(device.gpu_id, 0)


class TestInput(unittest.TestCase):

def _verify_correctness(self, struct: torchtrt.Input, target: Dict) -> bool:
internal = struct._to_internal()

list_eq = lambda al, bl: all([a == b for (a, b) in zip (al, bl)])
list_eq = lambda al, bl: all([a == b for (a, b) in zip(al, bl)])

eq = lambda a, b : a == b
eq = lambda a, b: a == b

def field_is_correct(field, equal_fn, a1, a2):
equal = equal_fn(a1, a2)
Expand All @@ -371,12 +388,12 @@ def field_is_correct(field, equal_fn, a1, a2):
opt_ = field_is_correct("opt", list_eq, internal.opt, target["opt"])
max_ = field_is_correct("max", list_eq, internal.max, target["max"])
is_dynamic_ = field_is_correct("is_dynamic", eq, internal.input_is_dynamic, target["input_is_dynamic"])
explicit_set_dtype_ = field_is_correct("explicit_dtype", eq, internal._explicit_set_dtype, target["explicit_set_dtype"])
explicit_set_dtype_ = field_is_correct("explicit_dtype", eq, internal._explicit_set_dtype,
target["explicit_set_dtype"])
dtype_ = field_is_correct("dtype", eq, int(internal.dtype), int(target["dtype"]))
format_ = field_is_correct("format", eq, int(internal.format), int(target["format"]))

return all([min_,opt_,max_,is_dynamic_,explicit_set_dtype_,dtype_,format_])

return all([min_, opt_, max_, is_dynamic_, explicit_set_dtype_, dtype_, format_])

def test_infer_from_example_tensor(self):
shape = [1, 3, 255, 255]
Expand All @@ -394,7 +411,6 @@ def test_infer_from_example_tensor(self):
i = torchtrt.Input._from_tensor(example_tensor)
self.assertTrue(self._verify_correctness(i, target))


def test_static_shape(self):
shape = [1, 3, 255, 255]
target = {
Expand Down Expand Up @@ -482,9 +498,12 @@ def test_dynamic_shape(self):
self.assertTrue(self._verify_correctness(i, target))

tensor_shape = lambda shape: torch.randn(shape).shape
i = torchtrt.Input(min_shape=tensor_shape(min_shape), opt_shape=tensor_shape(opt_shape), max_shape=tensor_shape(max_shape))
i = torchtrt.Input(min_shape=tensor_shape(min_shape),
opt_shape=tensor_shape(opt_shape),
max_shape=tensor_shape(max_shape))
self.assertTrue(self._verify_correctness(i, target))


def test_suite():
suite = unittest.TestSuite()
suite.addTest(unittest.makeSuite(TestLoggingAPIs))
Expand Down
11 changes: 6 additions & 5 deletions tests/py/test_ptq_dataloader_calibrator.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,11 +26,12 @@ def setUp(self):
batch_size=1,
shuffle=False,
num_workers=1)
self.calibrator = torchtrt.ptq.DataLoaderCalibrator(self.testing_dataloader,
cache_file='./calibration.cache',
use_cache=False,
algo_type=torchtrt.ptq.CalibrationAlgo.ENTROPY_CALIBRATION_2,
device=torch.device('cuda:0'))
self.calibrator = torchtrt.ptq.DataLoaderCalibrator(
self.testing_dataloader,
cache_file='./calibration.cache',
use_cache=False,
algo_type=torchtrt.ptq.CalibrationAlgo.ENTROPY_CALIBRATION_2,
device=torch.device('cuda:0'))

def compute_accuracy(self, testing_dataloader, model):
total = 0
Expand Down
34 changes: 18 additions & 16 deletions tests/py/test_ptq_to_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,25 +26,27 @@ def setUp(self):
batch_size=1,
shuffle=False,
num_workers=1)
self.calibrator = torchtrt.ptq.DataLoaderCalibrator(self.testing_dataloader,
cache_file='./calibration.cache',
use_cache=False,
algo_type=torchtrt.ptq.CalibrationAlgo.ENTROPY_CALIBRATION_2,
device=torch.device('cuda:0'))
self.calibrator = torchtrt.ptq.DataLoaderCalibrator(
self.testing_dataloader,
cache_file='./calibration.cache',
use_cache=False,
algo_type=torchtrt.ptq.CalibrationAlgo.ENTROPY_CALIBRATION_2,
device=torch.device('cuda:0'))

self.spec = {
"forward":
torchtrt.ts.TensorRTCompileSpec(**{
"inputs": [torchtrt.Input([1, 3, 32, 32])],
"enabled_precisions": {torch.float, torch.half, torch.int8},
"calibrator": self.calibrator,
"device": {
"device_type": torchtrt.DeviceType.GPU,
"gpu_id": 0,
"dla_core": 0,
"allow_gpu_fallback": False,
}
})
torchtrt.ts.TensorRTCompileSpec(
**{
"inputs": [torchtrt.Input([1, 3, 32, 32])],
"enabled_precisions": {torch.float, torch.half, torch.int8},
"calibrator": self.calibrator,
"device": {
"device_type": torchtrt.DeviceType.GPU,
"gpu_id": 0,
"dla_core": 0,
"allow_gpu_fallback": False,
}
})
}

def compute_accuracy(self, testing_dataloader, model):
Expand Down
37 changes: 19 additions & 18 deletions tests/py/test_to_backend_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,24 +13,25 @@ def setUp(self):
self.scripted_model = torch.jit.script(self.model)
self.spec = {
"forward":
torchtrt.ts.TensorRTCompileSpec(**{
"inputs": [torchtrt.Input([1, 3, 300, 300])],
"enabled_precisions": {torch.float},
"refit": False,
"debug": False,
"strict_types": False,
"device": {
"device_type": torchtrt.DeviceType.GPU,
"gpu_id": 0,
"dla_core": 0,
"allow_gpu_fallback": True
},
"capability": torchtrt.EngineCapability.default,
"num_min_timing_iters": 2,
"num_avg_timing_iters": 1,
"max_batch_size": 0,
"disable_tf32": False,
})
torchtrt.ts.TensorRTCompileSpec(
**{
"inputs": [torchtrt.Input([1, 3, 300, 300])],
"enabled_precisions": {torch.float},
"refit": False,
"debug": False,
"strict_types": False,
"device": {
"device_type": torchtrt.DeviceType.GPU,
"gpu_id": 0,
"dla_core": 0,
"allow_gpu_fallback": True
},
"capability": torchtrt.EngineCapability.default,
"num_min_timing_iters": 2,
"num_avg_timing_iters": 1,
"max_batch_size": 0,
"disable_tf32": False,
})
}

def test_to_backend_lowering(self):
Expand Down

0 comments on commit 8143489

Please sign in to comment.