Add custom onnxruntime

espnet · Jul 7, 2022 · 6e5cd10 · 6e5cd10
1 parent 87e2c1e
commit 6e5cd10
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
     "install": [
         "setuptools>=38.5.1",
         "librosa>=0.8.0",
-        "onnxruntime>=1.4.0",
+        "onnxruntime",
         "sentencepiece>=0.1.91,!=0.1.92",
         "typeguard>=2.7.0",
         "PyYAML>=5.1.2",

diff --git a/tests/test_optimize.py b/tests/test_optimize.py
@@ -13,20 +13,20 @@
 
 test_cases = [
     ['encoder', 'transformer', 4, 256, 3, 0, False],
-    ['encoder', 'contextual_block_transformer', 4, 256, 3, 0, False],
-    # ['encoder', 'transformer', 4, 256, 4, 0, True],
-    # ['encoder', 'contextual_block_transformer', 4, 256, 4, 0, True],
-    # ['decoder', 'transformer', 4, 256, 4, 0, True],
-    # ['lm', 'transformer_pe', 4, 256, 4, 0, True],
-    # ['lm', 'transformer', 4, 256, 4, 0, True],
+    # ['encoder', 'contextual_block_transformer', 4, 256, 3, 0, False],
+    ['encoder', 'transformer', 4, 256, 3, 0, True],
+    # ['encoder', 'contextual_block_transformer', 4, 256, 3, 0, True],
+    ['decoder', 'transformer', 4, 256, 0, 6, True],
+    ['lm', 'transformer_pe', 4, 256, 0, 3, True],
+    ['lm', 'transformer', 4, 256, 0, 2, True],
 ]
 
 @pytest.mark.parametrize('model_type, model_name, n_head, h_size, n_att, n_cross_att, use_custom_ort', test_cases)
 def test_optimize(model_type, model_name, n_head, h_size, n_att, n_cross_att, use_custom_ort, model_export):
     export_dir = model_export.cache_dir / 'test' / \
         model_type / f'cache_{model_name}'
     output_dir = model_export.cache_dir / 'test' / \
-        'optimize' / f'cache_{model_name}'
+        'optimize' / model_type / f'cache_{model_name}'
 
     input_model = glob.glob(os.path.join(export_dir , f'*{model_type}*'))[0]
     model_name = os.path.basename(input_model)