def convert_to_onnx(self, onnx_output_dir=None, set_onnx_arg=True):
"""Convert the model to ONNX format and save to output_dir
Args:
onnx_output_dir (str, optional): If specified, ONNX model will be saved to output_dir (else args.output_dir will be used). Defaults to None.
set_onnx_arg (bool, optional): Updates the model args to set onnx=True. Defaults to True.
""" # noqa
if not onnx_output_dir:
onnx_output_dir = os.path.join(self.options.output_dir,
self.options.model_type,
self.options.model_name, "onnx")
os.makedirs(onnx_output_dir, exist_ok=True)
if not os.listdir(onnx_output_dir):
onnx_model_name = os.path.join(onnx_output_dir, "onnx_model.onnx")
with tempfile.TemporaryDirectory() as temp_dir:
basedir = os.path.basename(temp_dir)
temp_dir = os.path.join(self.options.output_dir, basedir)
self.save_model(output_dir=temp_dir, model=self.model)
convert(
framework="pt",
model=temp_dir,
tokenizer=self.tokenizer,
output=Path(onnx_model_name),
pipeline_name="ner",
opset=11,
)
self.tokenizer.save_pretrained(onnx_output_dir)
self.config.save_pretrained(onnx_output_dir)
onnx_options = SessionOptions()
use_cuda = True if self._device.type != 'cpu' else False
onnx_execution_provider = "CUDAExecutionProvider" if use_cuda else "CPUExecutionProvider"
onnx_options.intra_op_num_threads = 1
onnx_options.execution_mode = ExecutionMode.ORT_SEQUENTIAL
onnx_model_path = os.path.join(onnx_output_dir, "onnx_model.onnx")
if self.options.dynamic_quantize:
# Append "-quantized" at the end of the model's name
quantized_model_path = generate_identified_filename(
Path(onnx_model_path), "-quantized")
quantize_dynamic(Path(onnx_model_path), quantized_model_path)
onnx_model_path = quantized_model_path.as_posix()
return InferenceSession(onnx_model_path,
onnx_options,
providers=[onnx_execution_provider])
def test_generate_identified_name(self):
generated = generate_identified_filename(
Path("/home/something/my_fake_model.onnx"), "-test")
self.assertEqual("/home/something/my_fake_model-test.onnx",
generated.as_posix())