def _onnx_export(self,
test_name,
name,
model_name,
feature,
onnx_config_class_constructor,
device="cpu"):
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)
# 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
onnx_config = onnx_config_class_constructor(model.config)
if is_torch_available():
from transformers.utils import torch_version
if torch_version < onnx_config.torch_onnx_minimum_version:
pytest.skip(
"Skipping due to incompatible PyTorch version. Minimum required is"
f" {onnx_config.torch_onnx_minimum_version}, got: {torch_version}"
)
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
with NamedTemporaryFile("w") as output:
try:
onnx_inputs, onnx_outputs = export(
preprocessor,
model,
onnx_config,
onnx_config.default_onnx_opset,
Path(output.name),
device=device)
validate_model_outputs(
onnx_config,
preprocessor,
model,
Path(output.name),
onnx_outputs,
onnx_config.atol_for_validation,
)
except (RuntimeError, ValueError) as e:
self.fail(f"{name}, {feature} -> {e}")
def _onnx_export(self,
test_name,
name,
model_name,
feature,
onnx_config_class_constructor,
device="cpu"):
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():
from transformers.utils import torch_version
if torch_version < onnx_config.torch_onnx_minimum_version:
pytest.skip(
"Skipping due to incompatible PyTorch version. Minimum required is"
f" {onnx_config.torch_onnx_minimum_version}, got: {torch_version}"
)
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
with NamedTemporaryFile("w") as output:
try:
onnx_inputs, onnx_outputs = export(
preprocessor,
model,
onnx_config,
onnx_config.default_onnx_opset,
Path(output.name),
device=device)
validate_model_outputs(
onnx_config,
preprocessor,
model,
Path(output.name),
onnx_outputs,
onnx_config.atol_for_validation,
)
except (RuntimeError, ValueError) as e:
self.fail(f"{name}, {feature} -> {e}")