Unverified Commit b41cc0b8 authored by Funtowicz Morgan's avatar Funtowicz Morgan Committed by GitHub
Browse files

Fix flaky ONNX tests (#6531)

parent 39c3b1d9
import unittest import unittest
from os.path import dirname, exists
from pathlib import Path from pathlib import Path
from shutil import rmtree
from tempfile import NamedTemporaryFile, TemporaryDirectory from tempfile import NamedTemporaryFile, TemporaryDirectory
from transformers import BertConfig, BertTokenizerFast, FeatureExtractionPipeline from transformers import BertConfig, BertTokenizerFast, FeatureExtractionPipeline
...@@ -72,7 +70,7 @@ class OnnxExportTestCase(unittest.TestCase): ...@@ -72,7 +70,7 @@ class OnnxExportTestCase(unittest.TestCase):
def test_quantize_pytorch(self): def test_quantize_pytorch(self):
for model in OnnxExportTestCase.MODEL_TO_TEST: for model in OnnxExportTestCase.MODEL_TO_TEST:
path = self._test_export(model, "pt", 12) path = self._test_export(model, "pt", 12)
quantized_path = quantize(Path(path)) quantized_path = quantize(path)
# Ensure the actual quantized model is not bigger than the original one # Ensure the actual quantized model is not bigger than the original one
if quantized_path.stat().st_size >= Path(path).stat().st_size: if quantized_path.stat().st_size >= Path(path).stat().st_size:
...@@ -82,16 +80,16 @@ class OnnxExportTestCase(unittest.TestCase): ...@@ -82,16 +80,16 @@ class OnnxExportTestCase(unittest.TestCase):
try: try:
# Compute path # Compute path
with TemporaryDirectory() as tempdir: with TemporaryDirectory() as tempdir:
path = tempdir + "/model.onnx" path = Path(tempdir).joinpath("model.onnx")
# Remove folder if exists # Remove folder if exists
if exists(dirname(path)): if path.parent.exists():
rmtree(dirname(path)) path.parent.rmdir()
# Export # Export
convert(framework, model, path, opset, tokenizer) convert(framework, model, path, opset, tokenizer)
return path return path
except Exception as e: except Exception as e:
self.fail(e) self.fail(e)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment