Skip to content

Commit 8143489

Browse files
narendasannarendasan
authored
Merge pull request #700 from NVIDIA/test_global_compile
test: Add a test for the global compile function
2 parents 0b316d2 + 19b448c commit 8143489

File tree

5 files changed

+118
-69
lines changed

5 files changed

+118
-69
lines changed

‎.github/workflows/linter.yml

+30-4
Original file line numberDiff line numberDiff line change
@@ -8,14 +8,27 @@ jobs:
88
cpp-linting:
99
name: C++ Linting
1010
runs-on: ubuntu-latest
11+
permissions:
12+
actions: write
13+
checks: write
14+
contents: write
15+
deployments: none
16+
id-token: write
17+
issues: write
18+
discussions: write
19+
packages: write
20+
pull-requests: write
21+
repository-projects: none
22+
security-events: none
23+
statuses: write
1124
steps:
1225
- uses: actions/checkout@v2
1326
with:
1427
ref: ${{ github.event.pull_request.head.sha }}
1528
- name: Docker login
16-
uses: docker/#-action@f054a8b539a109f9f41c372932f1ae047eff08c9
29+
uses: docker/#-action@v1
1730
with:
18-
registry: ${{ env.REGISTRY }}
31+
registry: ghcr.io
1932
username: ${{ github.actor }}
2033
password: ${{ secrets.GITHUB_TOKEN }}
2134
- name: Run image
@@ -32,14 +45,27 @@ jobs:
3245
py-linting:
3346
name: Python Linting
3447
runs-on: ubuntu-latest
48+
permissions:
49+
actions: write
50+
checks: write
51+
contents: write
52+
deployments: none
53+
id-token: write
54+
issues: write
55+
discussions: write
56+
packages: write
57+
pull-requests: write
58+
repository-projects: none
59+
security-events: none
60+
statuses: write
3561
steps:
3662
- uses: actions/checkout@v2
3763
with:
3864
ref: ${{ github.event.pull_request.head.sha }}
3965
- name: Docker login
40-
uses: docker/#-action@f054a8b539a109f9f41c372932f1ae047eff08c9
66+
uses: docker/#-action@v1
4167
with:
42-
registry: ${{ env.REGISTRY }}
68+
registry: ghcr.io
4369
username: ${{ github.actor }}
4470
password: ${{ secrets.GITHUB_TOKEN }}
4571
- name: Run image

‎tests/py/test_api.py

+45-26
Original file line numberDiff line numberDiff line change
@@ -31,9 +31,25 @@ def test_compile_traced(self):
3131

3232
def test_compile_script(self):
3333
trt_mod = torchtrt.ts.compile(self.scripted_model,
34-
inputs=[self.input],
35-
device=torchtrt.Device(gpu_id=0),
36-
enabled_precisions={torch.float})
34+
inputs=[self.input],
35+
device=torchtrt.Device(gpu_id=0),
36+
enabled_precisions={torch.float})
37+
same = (trt_mod(self.input) - self.scripted_model(self.input)).abs().max()
38+
self.assertTrue(same < 2e-2)
39+
40+
def test_compile_global(self):
41+
trt_mod = torchtrt.compile(self.scripted_model,
42+
inputs=[self.input],
43+
device=torchtrt.Device(gpu_id=0),
44+
enabled_precisions={torch.float})
45+
same = (trt_mod(self.input) - self.scripted_model(self.input)).abs().max()
46+
self.assertTrue(same < 2e-2)
47+
48+
def test_compile_global_nn_mod(self):
49+
trt_mod = torchtrt.compile(self.model,
50+
inputs=[self.input],
51+
device=torchtrt.Device(gpu_id=0),
52+
enabled_precisions={torch.float})
3753
same = (trt_mod(self.input) - self.scripted_model(self.input)).abs().max()
3854
self.assertTrue(same < 2e-2)
3955

@@ -207,16 +223,16 @@ def setUp(self):
207223
def test_input_use_default_fp32(self):
208224
ts_model = torch.jit.script(self.model)
209225
trt_mod = torchtrt.ts.compile(ts_model,
210-
inputs=[torchtrt.Input(self.input.shape)],
211-
enabled_precisions={torch.float, torch.half})
226+
inputs=[torchtrt.Input(self.input.shape)],
227+
enabled_precisions={torch.float, torch.half})
212228
trt_mod(self.input)
213229

214230
def test_input_respect_user_setting_fp32_weights_fp16_in(self):
215231
ts_model = torch.jit.script(self.model)
216232
trt_mod = torchtrt.ts.compile(ts_model,
217-
inputs=[self.input.half()],
218-
require_full_compilation=True,
219-
enabled_precisions={torch.float, torch.half})
233+
inputs=[self.input.half()],
234+
require_full_compilation=True,
235+
enabled_precisions={torch.float, torch.half})
220236
trt_mod(self.input.half())
221237

222238
def test_input_respect_user_setting_fp32_weights_fp16_in_non_constructor(self):
@@ -225,9 +241,9 @@ def test_input_respect_user_setting_fp32_weights_fp16_in_non_constructor(self):
225241
input_spec.dtype = torch.half
226242

227243
trt_mod = torchtrt.ts.compile(ts_model,
228-
inputs=[input_spec],
229-
require_full_compilation=True,
230-
enabled_precisions={torch.float, torch.half})
244+
inputs=[input_spec],
245+
require_full_compilation=True,
246+
enabled_precisions={torch.float, torch.half})
231247
trt_mod(self.input.half())
232248

233249

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

243259
trt_mod = torchtrt.ts.compile(half_mod,
244-
inputs=[torchtrt.Input(self.input.shape)],
245-
enabled_precisions={torch.float, torch.half})
260+
inputs=[torchtrt.Input(self.input.shape)],
261+
enabled_precisions={torch.float, torch.half})
246262
trt_mod(self.input.half())
247263

248264
def test_input_use_default_fp16_without_fp16_enabled(self):
@@ -257,9 +273,9 @@ def test_input_respect_user_setting_fp16_weights_fp32_in(self):
257273
half_mod.half()
258274

259275
trt_mod = torchtrt.ts.compile(half_mod,
260-
inputs=[self.input],
261-
require_full_compilation=True,
262-
enabled_precisions={torch.float, torch.half})
276+
inputs=[self.input],
277+
require_full_compilation=True,
278+
enabled_precisions={torch.float, torch.half})
263279
trt_mod(self.input)
264280

265281
def test_input_respect_user_setting_fp16_weights_fp32_in_non_constuctor(self):
@@ -270,9 +286,9 @@ def test_input_respect_user_setting_fp16_weights_fp32_in_non_constuctor(self):
270286
input_spec.dtype = torch.float
271287

272288
trt_mod = torchtrt.ts.compile(half_mod,
273-
inputs=[input_spec],
274-
require_full_compilation=True,
275-
enabled_precisions={torch.float, torch.half})
289+
inputs=[input_spec],
290+
require_full_compilation=True,
291+
enabled_precisions={torch.float, torch.half})
276292
trt_mod(self.input)
277293

278294

@@ -352,14 +368,15 @@ def test_from_torch(self):
352368
self.assertEqual(device.device_type, torchtrt.DeviceType.GPU)
353369
self.assertEqual(device.gpu_id, 0)
354370

371+
355372
class TestInput(unittest.TestCase):
356373

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

360-
list_eq = lambda al, bl: all([a == b for (a, b) in zip (al, bl)])
377+
list_eq = lambda al, bl: all([a == b for (a, b) in zip(al, bl)])
361378

362-
eq = lambda a, b : a == b
379+
eq = lambda a, b: a == b
363380

364381
def field_is_correct(field, equal_fn, a1, a2):
365382
equal = equal_fn(a1, a2)
@@ -371,12 +388,12 @@ def field_is_correct(field, equal_fn, a1, a2):
371388
opt_ = field_is_correct("opt", list_eq, internal.opt, target["opt"])
372389
max_ = field_is_correct("max", list_eq, internal.max, target["max"])
373390
is_dynamic_ = field_is_correct("is_dynamic", eq, internal.input_is_dynamic, target["input_is_dynamic"])
374-
explicit_set_dtype_ = field_is_correct("explicit_dtype", eq, internal._explicit_set_dtype, target["explicit_set_dtype"])
391+
explicit_set_dtype_ = field_is_correct("explicit_dtype", eq, internal._explicit_set_dtype,
392+
target["explicit_set_dtype"])
375393
dtype_ = field_is_correct("dtype", eq, int(internal.dtype), int(target["dtype"]))
376394
format_ = field_is_correct("format", eq, int(internal.format), int(target["format"]))
377395

378-
return all([min_,opt_,max_,is_dynamic_,explicit_set_dtype_,dtype_,format_])
379-
396+
return all([min_, opt_, max_, is_dynamic_, explicit_set_dtype_, dtype_, format_])
380397

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

397-
398414
def test_static_shape(self):
399415
shape = [1, 3, 255, 255]
400416
target = {
@@ -482,9 +498,12 @@ def test_dynamic_shape(self):
482498
self.assertTrue(self._verify_correctness(i, target))
483499

484500
tensor_shape = lambda shape: torch.randn(shape).shape
485-
i = torchtrt.Input(min_shape=tensor_shape(min_shape), opt_shape=tensor_shape(opt_shape), max_shape=tensor_shape(max_shape))
501+
i = torchtrt.Input(min_shape=tensor_shape(min_shape),
502+
opt_shape=tensor_shape(opt_shape),
503+
max_shape=tensor_shape(max_shape))
486504
self.assertTrue(self._verify_correctness(i, target))
487505

506+
488507
def test_suite():
489508
suite = unittest.TestSuite()
490509
suite.addTest(unittest.makeSuite(TestLoggingAPIs))

‎tests/py/test_ptq_dataloader_calibrator.py

+6-5
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,12 @@ def setUp(self):
2626
batch_size=1,
2727
shuffle=False,
2828
num_workers=1)
29-
self.calibrator = torchtrt.ptq.DataLoaderCalibrator(self.testing_dataloader,
30-
cache_file='./calibration.cache',
31-
use_cache=False,
32-
algo_type=torchtrt.ptq.CalibrationAlgo.ENTROPY_CALIBRATION_2,
33-
device=torch.device('cuda:0'))
29+
self.calibrator = torchtrt.ptq.DataLoaderCalibrator(
30+
self.testing_dataloader,
31+
cache_file='./calibration.cache',
32+
use_cache=False,
33+
algo_type=torchtrt.ptq.CalibrationAlgo.ENTROPY_CALIBRATION_2,
34+
device=torch.device('cuda:0'))
3435

3536
def compute_accuracy(self, testing_dataloader, model):
3637
total = 0

‎tests/py/test_ptq_to_backend.py

+18-16
Original file line numberDiff line numberDiff line change
@@ -26,25 +26,27 @@ def setUp(self):
2626
batch_size=1,
2727
shuffle=False,
2828
num_workers=1)
29-
self.calibrator = torchtrt.ptq.DataLoaderCalibrator(self.testing_dataloader,
30-
cache_file='./calibration.cache',
31-
use_cache=False,
32-
algo_type=torchtrt.ptq.CalibrationAlgo.ENTROPY_CALIBRATION_2,
33-
device=torch.device('cuda:0'))
29+
self.calibrator = torchtrt.ptq.DataLoaderCalibrator(
30+
self.testing_dataloader,
31+
cache_file='./calibration.cache',
32+
use_cache=False,
33+
algo_type=torchtrt.ptq.CalibrationAlgo.ENTROPY_CALIBRATION_2,
34+
device=torch.device('cuda:0'))
3435

3536
self.spec = {
3637
"forward":
37-
torchtrt.ts.TensorRTCompileSpec(**{
38-
"inputs": [torchtrt.Input([1, 3, 32, 32])],
39-
"enabled_precisions": {torch.float, torch.half, torch.int8},
40-
"calibrator": self.calibrator,
41-
"device": {
42-
"device_type": torchtrt.DeviceType.GPU,
43-
"gpu_id": 0,
44-
"dla_core": 0,
45-
"allow_gpu_fallback": False,
46-
}
47-
})
38+
torchtrt.ts.TensorRTCompileSpec(
39+
**{
40+
"inputs": [torchtrt.Input([1, 3, 32, 32])],
41+
"enabled_precisions": {torch.float, torch.half, torch.int8},
42+
"calibrator": self.calibrator,
43+
"device": {
44+
"device_type": torchtrt.DeviceType.GPU,
45+
"gpu_id": 0,
46+
"dla_core": 0,
47+
"allow_gpu_fallback": False,
48+
}
49+
})
4850
}
4951

5052
def compute_accuracy(self, testing_dataloader, model):

‎tests/py/test_to_backend_api.py

+19-18
Original file line numberDiff line numberDiff line change
@@ -13,24 +13,25 @@ def setUp(self):
1313
self.scripted_model = torch.jit.script(self.model)
1414
self.spec = {
1515
"forward":
16-
torchtrt.ts.TensorRTCompileSpec(**{
17-
"inputs": [torchtrt.Input([1, 3, 300, 300])],
18-
"enabled_precisions": {torch.float},
19-
"refit": False,
20-
"debug": False,
21-
"strict_types": False,
22-
"device": {
23-
"device_type": torchtrt.DeviceType.GPU,
24-
"gpu_id": 0,
25-
"dla_core": 0,
26-
"allow_gpu_fallback": True
27-
},
28-
"capability": torchtrt.EngineCapability.default,
29-
"num_min_timing_iters": 2,
30-
"num_avg_timing_iters": 1,
31-
"max_batch_size": 0,
32-
"disable_tf32": False,
33-
})
16+
torchtrt.ts.TensorRTCompileSpec(
17+
**{
18+
"inputs": [torchtrt.Input([1, 3, 300, 300])],
19+
"enabled_precisions": {torch.float},
20+
"refit": False,
21+
"debug": False,
22+
"strict_types": False,
23+
"device": {
24+
"device_type": torchtrt.DeviceType.GPU,
25+
"gpu_id": 0,
26+
"dla_core": 0,
27+
"allow_gpu_fallback": True
28+
},
29+
"capability": torchtrt.EngineCapability.default,
30+
"num_min_timing_iters": 2,
31+
"num_avg_timing_iters": 1,
32+
"max_batch_size": 0,
33+
"disable_tf32": False,
34+
})
3435
}
3536

3637
def test_to_backend_lowering(self):

0 commit comments

Comments
 (0)