Error in converting Azure AutoML model to ONNX format

Question

I am working with Azure AutoML and have created a trial using different models. In the pipeline, the final step is a PipelineWithYTransformations.

Now, I am attempting to convert the generated .pkl model to an ONNX model, but I encounter the following error:

raise RuntimeError("Unable to find a shape calculator for type "
RuntimeError: Unable to find a shape calculator for type '<class 'azureml.training.tabular.models.pipeline_with_ytransformations.PipelineWithYTransformations'>'.

I am unsure how to handle this error and successfully convert the model to ONNX format.

Here are some additional details about my setup:

• The AutoML trial was created with the goal of training a machine learning model.

from azure.ai.ml import automl, Input
from azure.ai.ml.constants import AssetTypes

dataset_name = "DataSetName"
my_training_data_input = Input(type=AssetTypes.MLTABLE, path=dataset_name)


# configure the classification job
classification_job = automl.classification(
    compute="my_compute",
    experiment_name="my_experiment",
    training_data=my_training_data_input,
    target_column_name="Label",
    primary_metric="accuracy",
    n_cross_validations=5,
    enable_model_explainability=True,
    tags={"my_custom_tag": "My custom value"}
)

# Limits are all optional
classification_job.set_limits(
    timeout_minutes=600, 
    trial_timeout_minutes=20, 
    max_trials=2,
    enable_early_termination=True,
)

# Training properties are optional
classification_job.set_training(
    blocked_training_algorithms=["logistic_regression"], 
    enable_onnx_compatible_models=True
)

# Submit automl job
returned_job = ml_client.jobs.create_or_update(
    classification_job
)  

• I have tried using the convert function from the skl2onnx library to perform the conversion.
• The error seems to be related to the specific type of model (PipelineWithYTransformations) generated by AutoML. I would appreciate any guidance or suggestions on how to resolve this issue and successfully convert the model to ONNX format.

Thank you!

Answer 1

One possible solution to get ONNX format model with Azure Auto ML from your run object.

remote_run = experiment.submit(automl_config, show_output=False)

Is by using get_output:

best_run, onnx_mdl = remote_run.get_output(return_onnx_model=True)

Also, you can use OnnxConverter Class for saving the model in onnx format.

from azureml.automl.runtime.onnx_convert import OnnxConverter
onnx_fl_path = "./best_model.onnx"
OnnxConverter.save_onnx_model(onnx_mdl, onnx_fl_path)

For detailed steps and sample code please check this sample notebook.

For SDKv2 (Update) : On executing the provided script, the onnx format model will be generated alongside pickle model. To get the model you can use the below code snippet:

from mlflow.tracking.client import MlflowClient
from mlflow.artifacts import download_artifacts

# Initialize MLFlow client
mlflow_client = MlflowClient()

job_name = returned_job.name
mlflow_parent_run = mlflow_client.get_run(job_name)
# Get the best model's child run

best_child_run_id = mlflow_parent_run.data.tags["automl_best_child_run_id"]
print("Found best child run id: ", best_child_run_id)

best_run = mlflow_client.get_run(best_child_run_id)

print("Best child run: ")
print(best_run)
import os

# Create local folder
local_dir = "./artifact_downloads"
if not os.path.exists(local_dir):
    os.mkdir(local_dir)
    
# Download run's artifacts/outputs
local_path = download_artifacts(
    run_id=best_run.info.run_id, artifact_path="outputs", dst_path=local_dir
)
print("Artifacts downloaded in: {}".format(local_path))
print("Artifacts: {}".format(os.listdir(local_path)))

This will download all the artifacts including model.onnx file.

For more details, you can check this sample notebook.