support ONNX export of XDropout in deberta{,_v2} and sew_d (#17502)

* support ONNX export of XDropout in deberta{,_v2}

* black

* copy to sew_d

* add test

* isort

* use pytest.mark.filterwarnings

* review comments
This commit is contained in:
Gary Miguel
2022-08-03 03:33:44 -07:00
committed by GitHub
parent 92915ebec2
commit 9d7b70bcd7
4 changed files with 88 additions and 0 deletions

View File

@@ -1,3 +1,4 @@
import os
from pathlib import Path
from tempfile import NamedTemporaryFile
from unittest import TestCase
@@ -26,6 +27,11 @@ from transformers.testing_utils import require_onnx, require_rjieba, require_tf,
if is_torch_available() or is_tf_available():
from transformers.onnx.features import FeaturesManager
if is_torch_available():
import torch
from transformers.models.deberta import modeling_deberta
@require_onnx
class OnnxUtilsTestCaseV2(TestCase):
@@ -356,3 +362,40 @@ class OnnxExportTestCaseV2(TestCase):
self, test_name, name, model_name, feature, onnx_config_class_constructor
):
self._onnx_export(test_name, name, model_name, feature, onnx_config_class_constructor)
class StableDropoutTestCase(TestCase):
"""Tests export of StableDropout module."""
@require_torch
@pytest.mark.filterwarnings("ignore:.*Dropout.*:UserWarning:torch.onnx.*") # torch.onnx is spammy.
def test_training(self):
"""Tests export of StableDropout in training mode."""
devnull = open(os.devnull, "wb")
# drop_prob must be > 0 for the test to be meaningful
sd = modeling_deberta.StableDropout(0.1)
# Avoid warnings in training mode
do_constant_folding = False
# Dropout is a no-op in inference mode
training = torch.onnx.TrainingMode.PRESERVE
input = (torch.randn(2, 2),)
torch.onnx.export(
sd,
input,
devnull,
opset_version=12, # Minimum supported
do_constant_folding=do_constant_folding,
training=training,
)
# Expected to fail with opset_version < 12
with self.assertRaises(Exception):
torch.onnx.export(
sd,
input,
devnull,
opset_version=11,
do_constant_folding=do_constant_folding,
training=training,
)