def test_backbone_class_export(tmp_path):
stages_output = "tri_stage"
x = torch.randn(1, 3, 224, 224)
model = Backbone("resnet18", stages_output=stages_output, freeze=True).eval()
y = model(x)
assert tuple(o.size(1) for o in y) == model.stages_channel
## onnx
onnx_path = tmp_path.joinpath("resnet18_backbone.onnx")
onnx_export(model, x, str(onnx_path), input_names=['input'], opset_version=11)
sess = onnxruntime.InferenceSession(str(onnx_path), providers=['CPUExecutionProvider'])
input_name = sess.get_inputs()[0].name
ort_y = sess.run(None, {input_name: x.clone().numpy()})
assert tuple(o.shape[1] for o in ort_y) == model.stages_channel
assert all([np.allclose(r, p.numpy(), atol=1e-4) for r,p in zip(ort_y, y)])
pt_path = onnx_path.with_suffix(".pt")
pt_model = torchscript_trace(model, x, check_tolerance=1e-6)
pt_model.save(str(pt_path))
loaded_pt_model = torchscript_load(str(pt_path))
pt_y = loaded_pt_model(x.clone())
assert tuple(o.size(1) for o in pt_y) == model.stages_channel
assert all([torch.allclose(p, pr) for p,pr in zip(y, pt_y)])
def export_pytorch(
preprocessor: Union["PreTrainedTokenizer", "FeatureExtractionMixin"],
model: "PreTrainedModel",
config: OnnxConfig,
opset: int,
output: Path,
tokenizer: "PreTrainedTokenizer" = None,
device: str = "cpu",
) -> Tuple[List[str], List[str]]:
"""
Export a PyTorch model to an ONNX Intermediate Representation (IR)
Args:
preprocessor: ([`PreTrainedTokenizer`] or [`FeatureExtractionMixin`]):
The preprocessor used for encoding the data.
model ([`PreTrainedModel`]):
The model to export.
config ([`~onnx.config.OnnxConfig`]):
The ONNX configuration associated with the exported model.
opset (`int`):
The version of the ONNX operator set to use.
output (`Path`):
Directory to store the exported ONNX model.
device (`str`, *optional*, defaults to `cpu`):
The device on which the ONNX model will be exported. Either `cpu` or `cuda`.
Returns:
`Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from
the ONNX configuration.
"""
if isinstance(preprocessor, PreTrainedTokenizerBase) and tokenizer is not None:
raise ValueError("You cannot provide both a tokenizer and a preprocessor to export the model.")
if tokenizer is not None:
warnings.warn(
"The `tokenizer` argument is deprecated and will be removed in version 5 of Transformers. Use"
" `preprocessor` instead.",
FutureWarning,
)
logger.info("Overwriting the `preprocessor` argument with `tokenizer` to generate dummmy inputs.")
preprocessor = tokenizer
if issubclass(type(model), PreTrainedModel):
import torch
from torch.onnx import export as onnx_export
logger.info(f"Using framework PyTorch: {torch.__version__}")
with torch.no_grad():
model.config.return_dict = True
model.eval()
# Check if we need to override certain configuration item
if config.values_override is not None:
logger.info(f"Overriding {len(config.values_override)} configuration item(s)")
for override_config_key, override_config_value in config.values_override.items():
logger.info(f"\t- {override_config_key} -> {override_config_value}")
setattr(model.config, override_config_key, override_config_value)
# Ensure inputs match
# TODO: Check when exporting QA we provide "is_pair=True"
model_inputs = config.generate_dummy_inputs(preprocessor, framework=TensorType.PYTORCH)
device = torch.device(device)
if device.type == "cuda" and torch.cuda.is_available():
model.to(device)
model_inputs = dict((k, v.to(device)) for k, v in model_inputs.items())
inputs_match, matched_inputs = ensure_model_and_config_inputs_match(model, model_inputs.keys())
onnx_outputs = list(config.outputs.keys())
if not inputs_match:
raise ValueError("Model and config inputs doesn't match")
config.patch_ops()
# PyTorch deprecated the `enable_onnx_checker` and `use_external_data_format` arguments in v1.11,
# so we check the torch version for backwards compatibility
if parse(torch.__version__) < parse("1.10"):
# export can work with named args but the dict containing named args
# has to be the last element of the args tuple.
try:
onnx_export(
model,
(model_inputs,),
f=output.as_posix(),
input_names=list(config.inputs.keys()),
output_names=onnx_outputs,
dynamic_axes={
name: axes for name, axes in chain(config.inputs.items(), config.outputs.items())
},
do_constant_folding=True,
use_external_data_format=config.use_external_data_format(model.num_parameters()),
enable_onnx_checker=True,
opset_version=opset,
)
except RuntimeError as err:
message = str(err)
if (
message
== "Exporting model exceed maximum protobuf size of 2GB. Please call torch.onnx.export without"
" setting use_external_data_format parameter."
):
message = (
"Exporting model exceed maximum protobuf size of 2GB. Please call torch.onnx.export"
" without setting use_external_data_format parameter or try with torch 1.10+."
)
raise RuntimeError(message)
else:
raise err
else:
onnx_export(
model,
(model_inputs,),
f=output.as_posix(),
input_names=list(config.inputs.keys()),
output_names=onnx_outputs,
dynamic_axes={name: axes for name, axes in chain(config.inputs.items(), config.outputs.items())},
do_constant_folding=True,
opset_version=opset,
)
config.restore_ops()
return matched_inputs, onnx_outputs
def export_pytorch(
tokenizer: PreTrainedTokenizer,
model: PreTrainedModel,
config: OnnxConfig,
opset: int,
output: Path,
) -> Tuple[List[str], List[str]]:
"""
Export a PyTorch model to an ONNX Intermediate Representation (IR)
Args:
tokenizer ([`PreTrainedTokenizer`]):
The tokenizer used for encoding the data.
model ([`PreTrainedModel`]):
The model to export.
config ([`~onnx.config.OnnxConfig`]):
The ONNX configuration associated with the exported model.
opset (`int`):
The version of the ONNX operator set to use.
output (`Path`):
Directory to store the exported ONNX model.
Returns:
`Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from
the ONNX configuration.
"""
if issubclass(type(model), PreTrainedModel):
import torch
from torch.onnx import export as onnx_export
logger.info(f"Using framework PyTorch: {torch.__version__}")
with torch.no_grad():
model.config.return_dict = True
model.eval()
# Check if we need to override certain configuration item
if config.values_override is not None:
logger.info(
f"Overriding {len(config.values_override)} configuration item(s)"
)
for override_config_key, override_config_value in config.values_override.items(
):
logger.info(
f"\t- {override_config_key} -> {override_config_value}"
)
setattr(model.config, override_config_key,
override_config_value)
# Ensure inputs match
# TODO: Check when exporting QA we provide "is_pair=True"
model_inputs = config.generate_dummy_inputs(
tokenizer, framework=TensorType.PYTORCH)
inputs_match, matched_inputs = ensure_model_and_config_inputs_match(
model, model_inputs.keys())
onnx_outputs = list(config.outputs.keys())
if not inputs_match:
raise ValueError("Model and config inputs doesn't match")
config.patch_ops()
# PyTorch deprecated the `enable_onnx_checker` and `use_external_data_format` arguments in v1.11,
# so we check the torch version for backwards compatibility
if parse(torch.__version__) <= parse("1.10.99"):
# export can work with named args but the dict containing named args
# has to be the last element of the args tuple.
onnx_export(
model,
(model_inputs, ),
f=output.as_posix(),
input_names=list(config.inputs.keys()),
output_names=onnx_outputs,
dynamic_axes={
name: axes
for name, axes in chain(config.inputs.items(),
config.outputs.items())
},
do_constant_folding=True,
use_external_data_format=config.use_external_data_format(
model.num_parameters()),
enable_onnx_checker=True,
opset_version=opset,
)
else:
onnx_export(
model,
(model_inputs, ),
f=output.as_posix(),
input_names=list(config.inputs.keys()),
output_names=onnx_outputs,
dynamic_axes={
name: axes
for name, axes in chain(config.inputs.items(),
config.outputs.items())
},
do_constant_folding=True,
opset_version=opset,
)
config.restore_ops()
return matched_inputs, onnx_outputs