test_onnx_v2.py 22.2 KB
Newer Older
1
import os
2
import unittest
3
4
5
6
7
from pathlib import Path
from tempfile import NamedTemporaryFile
from unittest import TestCase
from unittest.mock import patch

lewtun's avatar
lewtun committed
8
import pytest
9
from packaging import version
10
from parameterized import parameterized
11

12
from transformers import AutoConfig, PreTrainedTokenizerBase, is_tf_available, is_torch_available
13
14
15
from transformers.onnx import (
    EXTERNAL_DATA_FORMAT_SIZE_LIMIT,
    OnnxConfig,
lewtun's avatar
lewtun committed
16
    OnnxConfigWithPast,
17
18
19
    ParameterFormat,
    validate_model_outputs,
)
20
21
22
23
24
from transformers.onnx.utils import (
    compute_effective_axis_dimension,
    compute_serialized_parameters_size,
    get_preprocessor,
)
25
from transformers.testing_utils import require_onnx, require_rjieba, require_tf, require_torch, require_vision, slow
26
27


28
if is_torch_available() or is_tf_available():
29
30
    from transformers.onnx.features import FeaturesManager

31
32
33
34
35
if is_torch_available():
    import torch

    from transformers.models.deberta import modeling_deberta

36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80

@require_onnx
class OnnxUtilsTestCaseV2(TestCase):
    """
    Cover all the utilities involved to export ONNX models
    """

    def test_compute_effective_axis_dimension(self):
        """
        When exporting ONNX model with dynamic axis (batch or sequence) we set batch_size and/or sequence_length = -1.
        We cannot generate an effective tensor with axis dim == -1, so we trick by using some "fixed" values
        (> 1 to avoid ONNX squeezing the axis).

        This test ensure we are correctly replacing generated batch / sequence tensor with axis > 1
        """

        # Dynamic axis (batch, no token added by the tokenizer)
        self.assertEqual(compute_effective_axis_dimension(-1, fixed_dimension=2, num_token_to_add=0), 2)

        # Static axis (batch, no token added by the tokenizer)
        self.assertEqual(compute_effective_axis_dimension(0, fixed_dimension=2, num_token_to_add=0), 2)

        # Dynamic axis (sequence, token added by the tokenizer 2 (no pair))
        self.assertEqual(compute_effective_axis_dimension(0, fixed_dimension=8, num_token_to_add=2), 6)
        self.assertEqual(compute_effective_axis_dimension(0, fixed_dimension=8, num_token_to_add=2), 6)

        # Dynamic axis (sequence, token added by the tokenizer 3 (pair))
        self.assertEqual(compute_effective_axis_dimension(0, fixed_dimension=8, num_token_to_add=3), 5)
        self.assertEqual(compute_effective_axis_dimension(0, fixed_dimension=8, num_token_to_add=3), 5)

    def test_compute_parameters_serialized_size(self):
        """
        This test ensures we compute a "correct" approximation of the underlying storage requirement (size) for all the
        parameters for the specified parameter's dtype.
        """
        self.assertEqual(compute_serialized_parameters_size(2, ParameterFormat.Float), 2 * ParameterFormat.Float.size)

    def test_flatten_output_collection_property(self):
        """
        This test ensures we correctly flatten nested collection such as the one we use when returning past_keys.
        past_keys = Tuple[Tuple]

        ONNX exporter will export nested collections as ${collection_name}.${level_idx_0}.${level_idx_1}...${idx_n}
        """
        self.assertEqual(
81
            OnnxConfig.flatten_output_collection_property("past_key", [[0], [1], [2]]),
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
            {
                "past_key.0": 0,
                "past_key.1": 1,
                "past_key.2": 2,
            },
        )


class OnnxConfigTestCaseV2(TestCase):
    """
    Cover the test for models default.

    Default means no specific features is being enabled on the model.
    """

    @patch.multiple(OnnxConfig, __abstractmethods__=set())
    def test_use_external_data_format(self):
        """
        External data format is required only if the serialized size of the parameters if bigger than 2Gb
        """
        TWO_GB_LIMIT = EXTERNAL_DATA_FORMAT_SIZE_LIMIT

        # No parameters
        self.assertFalse(OnnxConfig.use_external_data_format(0))

        # Some parameters
        self.assertFalse(OnnxConfig.use_external_data_format(1))

        # Almost 2Gb parameters
        self.assertFalse(OnnxConfig.use_external_data_format((TWO_GB_LIMIT - 1) // ParameterFormat.Float.size))

        # Exactly 2Gb parameters
        self.assertTrue(OnnxConfig.use_external_data_format(TWO_GB_LIMIT))

        # More than 2Gb parameters
        self.assertTrue(OnnxConfig.use_external_data_format((TWO_GB_LIMIT + 1) // ParameterFormat.Float.size))


class OnnxConfigWithPastTestCaseV2(TestCase):
    """
    Cover the tests for model which have use_cache feature (i.e. "with_past" for ONNX)
    """

125
126
127
128
129
130
    SUPPORTED_WITH_PAST_CONFIGS = {}
    # SUPPORTED_WITH_PAST_CONFIGS = {
    #     ("BART", BartConfig),
    #     ("GPT2", GPT2Config),
    #     # ("T5", T5Config)
    # }
131
132
133
134
135
136
137
138
139

    @patch.multiple(OnnxConfigWithPast, __abstractmethods__=set())
    def test_use_past(self):
        """
        Ensure the use_past variable is correctly being set
        """
        for name, config in OnnxConfigWithPastTestCaseV2.SUPPORTED_WITH_PAST_CONFIGS:
            with self.subTest(name):
                self.assertFalse(
140
141
                    OnnxConfigWithPast.from_model_config(config()).use_past,
                    "OnnxConfigWithPast.from_model_config() should not use_past",
142
143
144
                )

                self.assertTrue(
145
146
                    OnnxConfigWithPast.with_past(config()).use_past,
                    "OnnxConfigWithPast.from_model_config() should use_past",
147
148
149
150
151
152
153
154
155
156
                )

    @patch.multiple(OnnxConfigWithPast, __abstractmethods__=set())
    def test_values_override(self):
        """
        Ensure the use_past variable correctly set the `use_cache` value in model's configuration
        """
        for name, config in OnnxConfigWithPastTestCaseV2.SUPPORTED_WITH_PAST_CONFIGS:
            with self.subTest(name):
                # without past
157
                onnx_config_default = OnnxConfigWithPast.from_model_config(config())
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
                self.assertIsNotNone(onnx_config_default.values_override, "values_override should not be None")
                self.assertIn("use_cache", onnx_config_default.values_override, "use_cache should be present")
                self.assertFalse(
                    onnx_config_default.values_override["use_cache"], "use_cache should be False if not using past"
                )

                # with past
                onnx_config_default = OnnxConfigWithPast.with_past(config())
                self.assertIsNotNone(onnx_config_default.values_override, "values_override should not be None")
                self.assertIn("use_cache", onnx_config_default.values_override, "use_cache should be present")
                self.assertTrue(
                    onnx_config_default.values_override["use_cache"], "use_cache should be False if not using past"
                )


173
PYTORCH_EXPORT_MODELS = {
174
175
176
177
    ("albert", "hf-internal-testing/tiny-random-AlbertModel"),
    ("bert", "hf-internal-testing/tiny-random-BertModel"),
    ("beit", "microsoft/beit-base-patch16-224"),
    ("big-bird", "hf-internal-testing/tiny-random-BigBirdModel"),
178
    ("camembert", "camembert-base"),
179
180
181
182
183
184
185
186
    ("clip", "hf-internal-testing/tiny-random-CLIPModel"),
    ("convbert", "hf-internal-testing/tiny-random-ConvBertModel"),
    ("codegen", "hf-internal-testing/tiny-random-CodeGenModel"),
    ("data2vec-text", "hf-internal-testing/tiny-random-Data2VecTextModel"),
    ("data2vec-vision", "facebook/data2vec-vision-base"),
    ("deberta", "hf-internal-testing/tiny-random-DebertaModel"),
    ("deberta-v2", "hf-internal-testing/tiny-random-DebertaV2Model"),
    ("deit", "facebook/deit-small-patch16-224"),
187
    ("convnext", "facebook/convnext-tiny-224"),
regisss's avatar
regisss committed
188
    ("detr", "facebook/detr-resnet-50"),
189
190
191
192
    ("distilbert", "hf-internal-testing/tiny-random-DistilBertModel"),
    ("electra", "hf-internal-testing/tiny-random-ElectraModel"),
    ("groupvit", "nvidia/groupvit-gcc-yfcc"),
    ("ibert", "kssteven/ibert-roberta-base"),
193
    ("imagegpt", "openai/imagegpt-small"),
194
195
196
197
198
    ("levit", "facebook/levit-128S"),
    ("layoutlm", "hf-internal-testing/tiny-random-LayoutLMModel"),
    ("layoutlmv3", "microsoft/layoutlmv3-base"),
    ("longformer", "allenai/longformer-base-4096"),
    ("mobilebert", "hf-internal-testing/tiny-random-MobileBertModel"),
199
    ("mobilenet_v1", "google/mobilenet_v1_0.75_192"),
200
    ("mobilenet_v2", "google/mobilenet_v2_0.35_96"),
201
    ("mobilevit", "apple/mobilevit-small"),
202
    ("owlvit", "google/owlvit-base-patch32"),
203
204
    ("perceiver", "hf-internal-testing/tiny-random-PerceiverModel", ("masked-lm", "sequence-classification")),
    ("perceiver", "hf-internal-testing/tiny-random-PerceiverModel", ("image-classification",)),
205
    ("poolformer", "sail/poolformer_s12"),
Erin's avatar
Erin committed
206
    ("rembert", "google/rembert"),
207
208
209
    ("resnet", "microsoft/resnet-50"),
    ("roberta", "hf-internal-testing/tiny-random-RobertaModel"),
    ("roformer", "hf-internal-testing/tiny-random-RoFormerModel"),
210
    ("segformer", "nvidia/segformer-b0-finetuned-ade-512-512"),
211
    ("squeezebert", "hf-internal-testing/tiny-random-SqueezeBertModel"),
212
    ("swin", "microsoft/swin-tiny-patch4-window7-224"),
213
214
    ("vit", "google/vit-base-patch16-224"),
    ("yolos", "hustvl/yolos-tiny"),
215
    ("whisper", "openai/whisper-tiny.en"),
216
217
    ("xlm", "hf-internal-testing/tiny-random-XLMModel"),
    ("xlm-roberta", "hf-internal-testing/tiny-random-XLMRobertaXLModel"),
218
219
}

220
221
222
223
PYTORCH_EXPORT_ENCODER_DECODER_MODELS = {
    ("vision-encoder-decoder", "nlpconnect/vit-gpt2-image-captioning"),
}

224
PYTORCH_EXPORT_WITH_PAST_MODELS = {
225
226
227
    ("bloom", "hf-internal-testing/tiny-random-BloomModel"),
    ("gpt2", "hf-internal-testing/tiny-random-GPT2Model"),
    ("gpt-neo", "hf-internal-testing/tiny-random-GPTNeoModel"),
228
229
230
}

PYTORCH_EXPORT_SEQ2SEQ_WITH_PAST_MODELS = {
231
232
233
234
235
    ("bart", "hf-internal-testing/tiny-random-BartModel"),
    ("bigbird-pegasus", "hf-internal-testing/tiny-random-BigBirdPegasusModel"),
    ("blenderbot-small", "facebook/blenderbot_small-90M"),
    ("blenderbot", "hf-internal-testing/tiny-random-BlenderbotModel"),
    ("longt5", "hf-internal-testing/tiny-random-LongT5Model"),
236
    ("marian", "Helsinki-NLP/opus-mt-en-de"),
237
    ("mbart", "sshleifer/tiny-mbart"),
238
    ("mt5", "google/mt5-base"),
239
240
    ("m2m-100", "hf-internal-testing/tiny-random-M2M100Model"),
    ("t5", "hf-internal-testing/tiny-random-T5Model"),
241
242
}

243
# TODO(lewtun): Include the same model types in `PYTORCH_EXPORT_MODELS` once TensorFlow has parity with the PyTorch model implementations.
244
245
TENSORFLOW_EXPORT_DEFAULT_MODELS = {
    ("albert", "hf-internal-testing/tiny-albert"),
246
    ("bert", "hf-internal-testing/tiny-random-BertModel"),
247
    ("camembert", "camembert-base"),
248
249
    ("distilbert", "hf-internal-testing/tiny-random-DistilBertModel"),
    ("roberta", "hf-internal-testing/tiny-random-RobertaModel"),
250
251
}

252
253
# TODO(lewtun): Include the same model types in `PYTORCH_EXPORT_WITH_PAST_MODELS` once TensorFlow has parity with the PyTorch model implementations.
TENSORFLOW_EXPORT_WITH_PAST_MODELS = {}
254

255
256
# TODO(lewtun): Include the same model types in `PYTORCH_EXPORT_SEQ2SEQ_WITH_PAST_MODELS` once TensorFlow has parity with the PyTorch model implementations.
TENSORFLOW_EXPORT_SEQ2SEQ_WITH_PAST_MODELS = {}
257

258
259
260

def _get_models_to_test(export_models_list):
    models_to_test = []
261
    if is_torch_available() or is_tf_available():
262
263
264
265
266
267
        for name, model, *features in export_models_list:
            if features:
                feature_config_mapping = {
                    feature: FeaturesManager.get_config(name, feature) for _ in features for feature in _
                }
            else:
268
269
270
271
272
273
                # pre-process the model names
                model_type = name.replace("_", "-")
                model_name = getattr(model, "name", "")
                feature_config_mapping = FeaturesManager.get_supported_features_for_model_type(
                    model_type, model_name=model_name
                )
274
275

            for feature, onnx_config_class_constructor in feature_config_mapping.items():
276
277
278
279
280
                models_to_test.append((f"{name}_{feature}", name, model, feature, onnx_config_class_constructor))
        return sorted(models_to_test)
    else:
        # Returning some dummy test that should not be ever called because of the @require_torch / @require_tf
        # decorators.
281
282
        # The reason for not returning an empty list is because parameterized.expand complains when it's empty.
        return [("dummy", "dummy", "dummy", "dummy", OnnxConfig.from_model_config)]
283
284
285
286
287
288
289


class OnnxExportTestCaseV2(TestCase):
    """
    Integration tests ensuring supported models are correctly exported
    """

290
291
292
    def _onnx_export(
        self, test_name, name, model_name, feature, onnx_config_class_constructor, device="cpu", framework="pt"
    ):
293
294
        from transformers.onnx import export

295
        model_class = FeaturesManager.get_model_class_for_feature(feature, framework=framework)
lewtun's avatar
lewtun committed
296
        config = AutoConfig.from_pretrained(model_name)
297
        model = model_class.from_config(config)
Yih-Dar's avatar
Yih-Dar committed
298
299
300
301
302
303

        # Dynamic axes aren't supported for YOLO-like models. This means they cannot be exported to ONNX on CUDA devices.
        # See: https://github.com/ultralytics/yolov5/pull/8378
        if model.__class__.__name__.startswith("Yolos") and device != "cpu":
            return

304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
        # ONNX inference fails with the following name, feature, framework parameterizations
        # See: https://github.com/huggingface/transformers/issues/19357
        if (name, feature, framework) in {
            ("deberta-v2", "question-answering", "pt"),
            ("deberta-v2", "multiple-choice", "pt"),
            ("roformer", "multiple-choice", "pt"),
            ("groupvit", "default", "pt"),
            ("perceiver", "masked-lm", "pt"),
            ("perceiver", "sequence-classification", "pt"),
            ("perceiver", "image-classification", "pt"),
            ("bert", "multiple-choice", "tf"),
            ("camembert", "multiple-choice", "tf"),
            ("roberta", "multiple-choice", "tf"),
        }:
            return

320
        onnx_config = onnx_config_class_constructor(model.config)
321

lewtun's avatar
lewtun committed
322
        if is_torch_available():
323
            from transformers.utils import get_torch_version
lewtun's avatar
lewtun committed
324

325
            if version.parse(get_torch_version()) < onnx_config.torch_onnx_minimum_version:
lewtun's avatar
lewtun committed
326
                pytest.skip(
Sylvain Gugger's avatar
Sylvain Gugger committed
327
                    "Skipping due to incompatible PyTorch version. Minimum required is"
328
                    f" {onnx_config.torch_onnx_minimum_version}, got: {get_torch_version()}"
lewtun's avatar
lewtun committed
329
330
                )

331
332
333
334
335
        preprocessor = get_preprocessor(model_name)

        # Useful for causal lm models that do not use pad tokens.
        if isinstance(preprocessor, PreTrainedTokenizerBase) and not getattr(config, "pad_token_id", None):
            config.pad_token_id = preprocessor.eos_token_id
lewtun's avatar
lewtun committed
336

337
338
339
        with NamedTemporaryFile("w") as output:
            try:
                onnx_inputs, onnx_outputs = export(
340
                    preprocessor, model, onnx_config, onnx_config.default_onnx_opset, Path(output.name), device=device
341
342
343
                )
                validate_model_outputs(
                    onnx_config,
lewtun's avatar
lewtun committed
344
                    preprocessor,
345
346
347
348
349
350
351
                    model,
                    Path(output.name),
                    onnx_outputs,
                    onnx_config.atol_for_validation,
                )
            except (RuntimeError, ValueError) as e:
                self.fail(f"{name}, {feature} -> {e}")
352

353
354
355
356
357
358
359
360
361
362
363
364
365
    def _onnx_export_encoder_decoder_models(
        self, test_name, name, model_name, feature, onnx_config_class_constructor, device="cpu"
    ):
        from transformers import AutoFeatureExtractor, AutoTokenizer
        from transformers.onnx import export

        model_class = FeaturesManager.get_model_class_for_feature(feature)
        config = AutoConfig.from_pretrained(model_name)
        model = model_class.from_config(config)

        onnx_config = onnx_config_class_constructor(model.config)

        if is_torch_available():
366
            from transformers.utils import get_torch_version
367

368
            if version.parse(get_torch_version()) < onnx_config.torch_onnx_minimum_version:
369
370
                pytest.skip(
                    "Skipping due to incompatible PyTorch version. Minimum required is"
371
                    f" {onnx_config.torch_onnx_minimum_version}, got: {get_torch_version()}"
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
                )

        encoder_model = model.get_encoder()
        decoder_model = model.get_decoder()

        encoder_onnx_config = onnx_config.get_encoder_config(encoder_model.config)
        decoder_onnx_config = onnx_config.get_decoder_config(encoder_model.config, decoder_model.config, feature)

        preprocessor = AutoFeatureExtractor.from_pretrained(model_name)

        onnx_opset = max(encoder_onnx_config.default_onnx_opset, decoder_onnx_config.default_onnx_opset)

        with NamedTemporaryFile("w") as encoder_output:
            onnx_inputs, onnx_outputs = export(
                preprocessor, encoder_model, encoder_onnx_config, onnx_opset, Path(encoder_output.name), device=device
            )
            validate_model_outputs(
                encoder_onnx_config,
                preprocessor,
                encoder_model,
                Path(encoder_output.name),
                onnx_outputs,
                encoder_onnx_config.atol_for_validation,
            )

        preprocessor = AutoTokenizer.from_pretrained(model_name)

        with NamedTemporaryFile("w") as decoder_output:
400
            _, onnx_outputs = export(
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
                preprocessor,
                decoder_model,
                decoder_onnx_config,
                onnx_config.default_onnx_opset,
                Path(decoder_output.name),
                device=device,
            )
            validate_model_outputs(
                decoder_onnx_config,
                preprocessor,
                decoder_model,
                Path(decoder_output.name),
                onnx_outputs,
                decoder_onnx_config.atol_for_validation,
            )

417
    @parameterized.expand(_get_models_to_test(PYTORCH_EXPORT_MODELS))
418
419
    @slow
    @require_torch
lewtun's avatar
lewtun committed
420
    @require_vision
421
    @require_rjieba
422
    def test_pytorch_export(self, test_name, name, model_name, feature, onnx_config_class_constructor):
423
        self._onnx_export(test_name, name, model_name, feature, onnx_config_class_constructor)
424

425
426
427
428
429
430
431
432
    @parameterized.expand(_get_models_to_test(PYTORCH_EXPORT_MODELS))
    @slow
    @require_torch
    @require_vision
    @require_rjieba
    def test_pytorch_export_on_cuda(self, test_name, name, model_name, feature, onnx_config_class_constructor):
        self._onnx_export(test_name, name, model_name, feature, onnx_config_class_constructor, device="cuda")

433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
    @parameterized.expand(_get_models_to_test(PYTORCH_EXPORT_ENCODER_DECODER_MODELS))
    @slow
    @require_torch
    @require_vision
    @require_rjieba
    def test_pytorch_export_encoder_decoder_models(
        self, test_name, name, model_name, feature, onnx_config_class_constructor
    ):
        self._onnx_export_encoder_decoder_models(test_name, name, model_name, feature, onnx_config_class_constructor)

    @parameterized.expand(_get_models_to_test(PYTORCH_EXPORT_ENCODER_DECODER_MODELS))
    @slow
    @require_torch
    @require_vision
    @require_rjieba
    def test_pytorch_export_encoder_decoder_models_on_cuda(
        self, test_name, name, model_name, feature, onnx_config_class_constructor
    ):
        self._onnx_export_encoder_decoder_models(
            test_name, name, model_name, feature, onnx_config_class_constructor, device="cuda"
        )

455
456
457
458
    @parameterized.expand(_get_models_to_test(PYTORCH_EXPORT_WITH_PAST_MODELS))
    @slow
    @require_torch
    def test_pytorch_export_with_past(self, test_name, name, model_name, feature, onnx_config_class_constructor):
459
        self._onnx_export(test_name, name, model_name, feature, onnx_config_class_constructor)
460

461
462
463
464
465
466
    @parameterized.expand(_get_models_to_test(PYTORCH_EXPORT_SEQ2SEQ_WITH_PAST_MODELS))
    @slow
    @require_torch
    def test_pytorch_export_seq2seq_with_past(
        self, test_name, name, model_name, feature, onnx_config_class_constructor
    ):
467
468
469
470
471
        self._onnx_export(test_name, name, model_name, feature, onnx_config_class_constructor)

    @parameterized.expand(_get_models_to_test(TENSORFLOW_EXPORT_DEFAULT_MODELS))
    @slow
    @require_tf
lewtun's avatar
lewtun committed
472
    @require_vision
473
    def test_tensorflow_export(self, test_name, name, model_name, feature, onnx_config_class_constructor):
474
        self._onnx_export(test_name, name, model_name, feature, onnx_config_class_constructor, framework="tf")
475

476
    @parameterized.expand(_get_models_to_test(TENSORFLOW_EXPORT_WITH_PAST_MODELS), skip_on_empty=True)
477
478
479
    @slow
    @require_tf
    def test_tensorflow_export_with_past(self, test_name, name, model_name, feature, onnx_config_class_constructor):
480
        self._onnx_export(test_name, name, model_name, feature, onnx_config_class_constructor, framework="tf")
481

482
    @parameterized.expand(_get_models_to_test(TENSORFLOW_EXPORT_SEQ2SEQ_WITH_PAST_MODELS), skip_on_empty=True)
483
484
485
486
487
    @slow
    @require_tf
    def test_tensorflow_export_seq2seq_with_past(
        self, test_name, name, model_name, feature, onnx_config_class_constructor
    ):
488
        self._onnx_export(test_name, name, model_name, feature, onnx_config_class_constructor, framework="tf")
489
490
491
492
493


class StableDropoutTestCase(TestCase):
    """Tests export of StableDropout module."""

494
    @unittest.skip("torch 2.0.0 gives `torch.onnx.errors.OnnxExporterError: Module onnx is not installed!`.")
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
    @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,
            )