mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
3c54ff6bcd
Currently, the ONNX exporter using torch.nn.Module as input can support FakeTensor because the ONNX model stores all initializers When using torch.export.ExportedProgram as input, the initializers are lifted as inputs. In order to execute the ONNX model, we need to pass a reference to the non-fake model to the ONNXProgram.adapt_torch_inputs_to_onnx API, so that initializers can be fetched from the model and fed to the ONNX model as input ps: https://github.com/pytorch/pytorch/issues/115461 will track the API revision for the cases where additional `model_with_state_dict` are required to produce complete ONNX files exported with fake support. This is also tracked by the umbrella fake tensor issue https://github.com/pytorch/pytorch/issues/105464 FYI @BowenBao Pull Request resolved: https://github.com/pytorch/pytorch/pull/114407 Approved by: https://github.com/BowenBao
418 lines
12 KiB
Python
418 lines
12 KiB
Python
# Owner(s): ["module: onnx"]
|
|
from __future__ import annotations
|
|
|
|
import functools
|
|
import os
|
|
import random
|
|
import sys
|
|
import unittest
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
import packaging.version
|
|
import pytest
|
|
|
|
import torch
|
|
from torch.autograd import function
|
|
from torch.onnx._internal import diagnostics
|
|
from torch.testing._internal import common_utils
|
|
|
|
pytorch_test_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
|
|
sys.path.insert(-1, pytorch_test_dir)
|
|
|
|
torch.set_default_tensor_type("torch.FloatTensor")
|
|
|
|
BATCH_SIZE = 2
|
|
|
|
RNN_BATCH_SIZE = 7
|
|
RNN_SEQUENCE_LENGTH = 11
|
|
RNN_INPUT_SIZE = 5
|
|
RNN_HIDDEN_SIZE = 3
|
|
|
|
|
|
def _skipper(condition, reason):
|
|
def decorator(f):
|
|
@functools.wraps(f)
|
|
def wrapper(*args, **kwargs):
|
|
if condition():
|
|
raise unittest.SkipTest(reason)
|
|
return f(*args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return decorator
|
|
|
|
|
|
skipIfNoCuda = _skipper(lambda: not torch.cuda.is_available(), "CUDA is not available")
|
|
|
|
skipIfTravis = _skipper(lambda: os.getenv("TRAVIS"), "Skip In Travis")
|
|
|
|
skipIfNoBFloat16Cuda = _skipper(
|
|
lambda: not torch.cuda.is_bf16_supported(), "BFloat16 CUDA is not available"
|
|
)
|
|
|
|
skipIfQuantizationBackendQNNPack = _skipper(
|
|
lambda: torch.backends.quantized.engine == "qnnpack",
|
|
"Not compatible with QNNPack quantization backend",
|
|
)
|
|
|
|
|
|
# skips tests for all versions below min_opset_version.
|
|
# if exporting the op is only supported after a specific version,
|
|
# add this wrapper to prevent running the test for opset_versions
|
|
# smaller than the currently tested opset_version
|
|
def skipIfUnsupportedMinOpsetVersion(min_opset_version):
|
|
def skip_dec(func):
|
|
@functools.wraps(func)
|
|
def wrapper(self, *args, **kwargs):
|
|
if self.opset_version < min_opset_version:
|
|
raise unittest.SkipTest(
|
|
f"Unsupported opset_version: {self.opset_version} < {min_opset_version}"
|
|
)
|
|
return func(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return skip_dec
|
|
|
|
|
|
# skips tests for all versions above max_opset_version.
|
|
def skipIfUnsupportedMaxOpsetVersion(max_opset_version):
|
|
def skip_dec(func):
|
|
@functools.wraps(func)
|
|
def wrapper(self, *args, **kwargs):
|
|
if self.opset_version > max_opset_version:
|
|
raise unittest.SkipTest(
|
|
f"Unsupported opset_version: {self.opset_version} > {max_opset_version}"
|
|
)
|
|
return func(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return skip_dec
|
|
|
|
|
|
# skips tests for all opset versions.
|
|
def skipForAllOpsetVersions():
|
|
def skip_dec(func):
|
|
@functools.wraps(func)
|
|
def wrapper(self, *args, **kwargs):
|
|
if self.opset_version:
|
|
raise unittest.SkipTest(
|
|
"Skip verify test for unsupported opset_version"
|
|
)
|
|
return func(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return skip_dec
|
|
|
|
|
|
def skipTraceTest(skip_before_opset_version: Optional[int] = None, reason: str = ""):
|
|
"""Skip tracing test for opset version less than skip_before_opset_version.
|
|
|
|
Args:
|
|
skip_before_opset_version: The opset version before which to skip tracing test.
|
|
If None, tracing test is always skipped.
|
|
reason: The reason for skipping tracing test.
|
|
|
|
Returns:
|
|
A decorator for skipping tracing test.
|
|
"""
|
|
|
|
def skip_dec(func):
|
|
@functools.wraps(func)
|
|
def wrapper(self, *args, **kwargs):
|
|
if skip_before_opset_version is not None:
|
|
self.skip_this_opset = self.opset_version < skip_before_opset_version
|
|
else:
|
|
self.skip_this_opset = True
|
|
if self.skip_this_opset and not self.is_script:
|
|
raise unittest.SkipTest(f"Skip verify test for torch trace. {reason}")
|
|
return func(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return skip_dec
|
|
|
|
|
|
def skipScriptTest(skip_before_opset_version: Optional[int] = None, reason: str = ""):
|
|
"""Skip scripting test for opset version less than skip_before_opset_version.
|
|
|
|
Args:
|
|
skip_before_opset_version: The opset version before which to skip scripting test.
|
|
If None, scripting test is always skipped.
|
|
reason: The reason for skipping scripting test.
|
|
|
|
Returns:
|
|
A decorator for skipping scripting test.
|
|
"""
|
|
|
|
def skip_dec(func):
|
|
@functools.wraps(func)
|
|
def wrapper(self, *args, **kwargs):
|
|
if skip_before_opset_version is not None:
|
|
self.skip_this_opset = self.opset_version < skip_before_opset_version
|
|
else:
|
|
self.skip_this_opset = True
|
|
if self.skip_this_opset and self.is_script:
|
|
raise unittest.SkipTest(f"Skip verify test for TorchScript. {reason}")
|
|
return func(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return skip_dec
|
|
|
|
|
|
# NOTE: This decorator is currently unused, but we may want to use it in the future when
|
|
# we have more tests that are not supported in released ORT.
|
|
def skip_min_ort_version(reason: str, version: str, dynamic_only: bool = False):
|
|
def skip_dec(func):
|
|
@functools.wraps(func)
|
|
def wrapper(self, *args, **kwargs):
|
|
if (
|
|
packaging.version.parse(self.ort_version).release
|
|
< packaging.version.parse(version).release
|
|
):
|
|
if dynamic_only and not self.dynamic_shapes:
|
|
return func(self, *args, **kwargs)
|
|
|
|
raise unittest.SkipTest(
|
|
f"ONNX Runtime version: {version} is older than required version {version}. "
|
|
f"Reason: {reason}."
|
|
)
|
|
return func(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return skip_dec
|
|
|
|
|
|
def skip_dynamic_fx_test(reason: str, skip_model_type=None):
|
|
"""Skip dynamic exporting test.
|
|
|
|
Args:
|
|
reason: The reason for skipping dynamic exporting test.
|
|
skip_model_type (onnx_test_common.TorchModelType): The model type to skip dynamic exporting test for.
|
|
When None, model type is not used to skip dynamic tests.
|
|
|
|
Returns:
|
|
A decorator for skipping dynamic exporting test.
|
|
"""
|
|
|
|
def skip_dec(func):
|
|
@functools.wraps(func)
|
|
def wrapper(self, *args, **kwargs):
|
|
if self.dynamic_shapes and (
|
|
not skip_model_type or self.model_type == skip_model_type
|
|
):
|
|
raise unittest.SkipTest(
|
|
f"Skip verify dynamic shapes test for FX. {reason}"
|
|
)
|
|
return func(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return skip_dec
|
|
|
|
|
|
def skip_load_checkpoint_after_model_creation(reason: str):
|
|
"""Skip loading checkpoint right after model initialization.
|
|
|
|
Args:
|
|
reason: The reason for skipping dynamic exporting test.
|
|
|
|
Returns:
|
|
A decorator for skipping dynamic exporting test.
|
|
"""
|
|
|
|
def skip_dec(func):
|
|
@functools.wraps(func)
|
|
def wrapper(self, *args, **kwargs):
|
|
if self.load_checkpoint_during_init:
|
|
raise unittest.SkipTest(
|
|
f"Skip loading checkpoint during model initialization for FX tests. {reason}"
|
|
)
|
|
return func(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return skip_dec
|
|
|
|
|
|
def skip_op_level_debug_test(reason: str):
|
|
"""Skip tests with op_level_debug enabled.
|
|
|
|
Args:
|
|
reason: The reason for skipping tests with op_level_debug enabled.
|
|
|
|
Returns:
|
|
A decorator for skipping tests with op_level_debug enabled.
|
|
"""
|
|
|
|
def skip_dec(func):
|
|
@functools.wraps(func)
|
|
def wrapper(self, *args, **kwargs):
|
|
if self.op_level_debug:
|
|
raise unittest.SkipTest(
|
|
f"Skip test with op_level_debug enabled. {reason}"
|
|
)
|
|
return func(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return skip_dec
|
|
|
|
|
|
def skip_in_ci(reason: str):
|
|
"""Skip test in CI.
|
|
|
|
Args:
|
|
reason: The reason for skipping test in CI.
|
|
|
|
Returns:
|
|
A decorator for skipping test in CI.
|
|
"""
|
|
|
|
def skip_dec(func):
|
|
@functools.wraps(func)
|
|
def wrapper(self, *args, **kwargs):
|
|
if os.getenv("CI"):
|
|
raise unittest.SkipTest(f"Skip test in CI. {reason}")
|
|
return func(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return skip_dec
|
|
|
|
|
|
def xfail(reason: str):
|
|
"""Expect failure.
|
|
|
|
Args:
|
|
reason: The reason for expected failure.
|
|
|
|
Returns:
|
|
A decorator for expecting test failure.
|
|
"""
|
|
return unittest.expectedFailure
|
|
|
|
|
|
# skips tests for opset_versions listed in unsupported_opset_versions.
|
|
# if the caffe2 test cannot be run for a specific version, add this wrapper
|
|
# (for example, an op was modified but the change is not supported in caffe2)
|
|
def skipIfUnsupportedOpsetVersion(unsupported_opset_versions):
|
|
def skip_dec(func):
|
|
@functools.wraps(func)
|
|
def wrapper(self, *args, **kwargs):
|
|
if self.opset_version in unsupported_opset_versions:
|
|
raise unittest.SkipTest(
|
|
"Skip verify test for unsupported opset_version"
|
|
)
|
|
return func(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return skip_dec
|
|
|
|
|
|
def skipShapeChecking(func):
|
|
@functools.wraps(func)
|
|
def wrapper(self, *args, **kwargs):
|
|
self.check_shape = False
|
|
return func(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
|
|
def skipDtypeChecking(func):
|
|
@functools.wraps(func)
|
|
def wrapper(self, *args, **kwargs):
|
|
self.check_dtype = False
|
|
return func(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
|
|
def xfail_if_model_type_is_exportedprogram(reason: str):
|
|
"""xfail test with models using ExportedProgram as input.
|
|
|
|
Args:
|
|
reason: The reason for xfail the ONNX export test.
|
|
|
|
Returns:
|
|
A decorator for xfail tests.
|
|
"""
|
|
|
|
import onnx_test_common
|
|
|
|
def xfail_dec(func):
|
|
@functools.wraps(func)
|
|
def wrapper(self, *args, **kwargs):
|
|
if (
|
|
self.model_type
|
|
== onnx_test_common.TorchModelType.TORCH_EXPORT_EXPORTEDPROGRAM
|
|
):
|
|
pytest.xfail(
|
|
reason=f"Xfail model_type==torch.export.ExportedProgram. {reason}"
|
|
)
|
|
return func(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return xfail_dec
|
|
|
|
|
|
def xfail_if_model_type_is_not_exportedprogram(reason: str):
|
|
"""xfail test without models using ExportedProgram as input.
|
|
|
|
Args:
|
|
reason: The reason for xfail the ONNX export test.
|
|
|
|
Returns:
|
|
A decorator for xfail tests.
|
|
"""
|
|
|
|
import onnx_test_common
|
|
|
|
def xfail_dec(func):
|
|
@functools.wraps(func)
|
|
def wrapper(self, *args, **kwargs):
|
|
if (
|
|
self.model_type
|
|
!= onnx_test_common.TorchModelType.TORCH_EXPORT_EXPORTEDPROGRAM
|
|
):
|
|
pytest.xfail(
|
|
reason=f"Xfail model_type!=torch.export.ExportedProgram. {reason}"
|
|
)
|
|
return func(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return xfail_dec
|
|
|
|
|
|
def flatten(x):
|
|
return tuple(function._iter_filter(lambda o: isinstance(o, torch.Tensor))(x))
|
|
|
|
|
|
def set_rng_seed(seed):
|
|
torch.manual_seed(seed)
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
|
|
|
|
class ExportTestCase(common_utils.TestCase):
|
|
"""Test case for ONNX export.
|
|
|
|
Any test case that tests functionalities under torch.onnx should inherit from this class.
|
|
"""
|
|
|
|
def setUp(self):
|
|
super().setUp()
|
|
# TODO(#88264): Flaky test failures after changing seed.
|
|
set_rng_seed(0)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed_all(0)
|
|
diagnostics.engine.clear()
|