Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
chenpangpang
transformers
Commits
b41cc0b8
"vscode:/vscode.git/clone" did not exist on "8e28c327fc832f7c6acffa2870c27db01658cc43"
Unverified
Commit
b41cc0b8
authored
Aug 17, 2020
by
Funtowicz Morgan
Committed by
GitHub
Aug 17, 2020
Browse files
Fix flaky ONNX tests (#6531)
parent
39c3b1d9
Changes
1
Show whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
7 additions
and
9 deletions
+7
-9
tests/test_onnx.py
tests/test_onnx.py
+7
-9
No files found.
tests/test_onnx.py
View file @
b41cc0b8
import
unittest
from
os.path
import
dirname
,
exists
from
pathlib
import
Path
from
shutil
import
rmtree
from
tempfile
import
NamedTemporaryFile
,
TemporaryDirectory
from
transformers
import
BertConfig
,
BertTokenizerFast
,
FeatureExtractionPipeline
...
...
@@ -72,7 +70,7 @@ class OnnxExportTestCase(unittest.TestCase):
def
test_quantize_pytorch
(
self
):
for
model
in
OnnxExportTestCase
.
MODEL_TO_TEST
:
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
if
quantized_path
.
stat
().
st_size
>=
Path
(
path
).
stat
().
st_size
:
...
...
@@ -82,11 +80,11 @@ class OnnxExportTestCase(unittest.TestCase):
try
:
# Compute path
with
TemporaryDirectory
()
as
tempdir
:
path
=
tempdir
+
"/
model.onnx"
path
=
Path
(
tempdir
).
joinpath
(
"
model.onnx"
)
# Remove folder if exists
if
exists
(
dirname
(
path
)
):
rmtree
(
dirname
(
path
)
)
if
path
.
parent
.
exists
(
):
path
.
parent
.
rmdir
(
)
# Export
convert
(
framework
,
model
,
path
,
opset
,
tokenizer
)
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment