From 628bf5e3c9a18d8a6626ba4e2c56714a9ea75bfd Mon Sep 17 00:00:00 2001 From: Mikhail Zolotukhin Date: Wed, 9 Jan 2019 21:11:58 -0800 Subject: [PATCH] test_jit.py: Speedup EndToEnd tests by reducing workload size. (#15906) Summary: Currently these tests are taking most of the time in test_jit.py run, with the proposed changes the testing time is reduced by ~75%: ``` TestEndToEndHybridFrontendModels.test_neural_style: 203.360s -> 10.650s TestEndToEndHybridFrontendModels.test_snli: 422.315s -> 9.152s TestEndToEndHybridFrontendModels.test_super_resolution: 73.362s -> 19.185s time python test/test_jit.py (real): 13m50.828s -> 3m11.768s time python test/test_jit.py (user): 85m59.745s -> 13m18.135s time python test/test_jit.py (sys): 144m9.028s -> 25m58.019s ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/15906 Differential Revision: D13619659 Pulled By: ZolotukhinM fbshipit-source-id: 6c22d8740f8ddb865c3a0667af32653723383816 --- test/test_jit.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/test/test_jit.py b/test/test_jit.py index 6102742..8df8dba 100644 --- a/test/test_jit.py +++ b/test/test_jit.py @@ -9371,7 +9371,7 @@ class TestEndToEndHybridFrontendModels(JitTestCase): out = self.conv2d(out) return out - self.checkTrace(TransformerNet(), (torch.rand(5, 3, 64, 64),), export_import=check_export_import) + self.checkTrace(TransformerNet(), (torch.rand(5, 3, 16, 16),), export_import=check_export_import) def test_neural_style(self): self._test_neural_style(self, device='cpu') @@ -9523,7 +9523,7 @@ class TestEndToEndHybridFrontendModels(JitTestCase): d_embed = 100 d_proj = 300 dp_ratio = 0.0 # For deterministic testing TODO: change by fixing seed in checkTrace? - d_hidden = 300 + d_hidden = 30 birnn = True d_out = 300 fix_emb = True @@ -9531,8 +9531,8 @@ class TestEndToEndHybridFrontendModels(JitTestCase): n_layers = 2 n_cells = 4 # 2 * n_layers because birnn = True - premise = torch.LongTensor(48, 128).random_(0, 100).to(device) - hypothesis = torch.LongTensor(24, 128).random_(0, 100).to(device) + premise = torch.LongTensor(48, 64).random_(0, 100).to(device) + hypothesis = torch.LongTensor(24, 64).random_(0, 100).to(device) if quantized: snli = SNLIClassifier(Config()).cpu() @@ -9584,7 +9584,7 @@ class TestEndToEndHybridFrontendModels(JitTestCase): return x net = Net(upscale_factor=4).to(device) - self.checkTrace(net, (torch.rand(5, 1, 64, 64, device=device),), + self.checkTrace(net, (torch.rand(5, 1, 32, 32, device=device),), export_import=check_export_import) @skipIfRocm -- 2.7.4