Download notebook

PyTorch to ONNX

eIQ AI Toolkit uses the Olive (https://microsoft.github.io/Olive/) framework for model conversions. This framework applies passes, where each pass represents a single transformation of a model. Passes can also be chained together to perform multiple transformations in sequence.

This guide explains how to use the OnnxConversion pass to convert a PyTorch model into ONNX format. It demonstrates how to use the eIQ AI Toolkit API to:

• Load a PyTorch model into the application

• Run the conversion to ONNX

• Retrieve the converted model

This guide requires the eIQ AI Toolkit backend to be running. If you haven’t set it up yet, please refer to the following tutorial: eIQ AI Toolkit setup & launch

import requests
import sys
from pathlib import Path

# Set your eIQ AI Toolkit url:
AI_TOOLKIT_BACKEND_URL = "http://localhost:8000"

Load PyTorch model

Loading any model into the application involves two steps:

1. Specify the model metadata

2. Upload the model file

The metadata must include the model type you are uploading. For example, use pytorch for a PyTorch model or onnx for an ONNX model.

For PyTorch models, you also need to define the following parameters:

• Input names

• Input shapes

• Input types

• Output names

You’ll see an example later showing how to structure and upload this data correctly.

eIQ AI Toolkit accepts PyTorch models only as a .zip archive containing:

• The model weights (.pth or .pt)

• A model definition script (.py)

The script must define a load_model function that accepts the path to the model weights and returns a torch.nn.Module instance (your model).

1. Prepare PyTorch model

If you already have a model, update the path to point to its location. If you don’t have a model yet, set the path to the location where the model will be saved. (See the following sections for instructions on downloading a sample model.)

# Modify the path to your PyTorch model
model_path = Path("path_to_pytorch_model.zip")

Use the following script to download the example model:

Note: Skip this step if you already have your own model.

example_model_url = "https://eiq.nxp.com/training-materials/_misc/models/model.zip"
with open(model_path, "wb") as f:
    response = requests.get(
        url=example_model_url
    )
    f.write(response.content)

2. Specify metadata

Now let’s specify the model metadata. As mentioned earlier, metadata for a PyTorch model must include:

• Input names

• Input shapes

• Input types

• Output names

The structure can look like this:

{
  "model_type": "pytorch",
  "io_config": {
    "input_config": [
      {
        "name": "images",
        "shape": [1,1,49,10],
        "type": "float32"
      },
      {
        "name": "descriptions",
        "shape": [1, 50, 50],
        "type": "float32"
      }
    ],
    "output_config": [
      {"name": "prediction"}
    ]
  }
}

In this example, we define a PyTorch model with two inputs—images and descriptions—and one output, prediction.

Next, let’s implement and use it in code.

Note: The code below is prepared for the example model. If you’re using your own model, update the values accordingly.

model_name = "PyTorch model"

# Define inputs metadata
inputs = [
            {
               "name": "images",
               "shape": [1,1,49,10],
               "type": "float32"
            }
         ]

# Define outputs metadata
outputs = [{"name": "y"}]

MODELS_API_URL = f"{AI_TOOLKIT_BACKEND_URL}/models"

# Full model metadata
model_metadata = {
        "model_type": "pytorch",
        "io_config": {
            "input_config": inputs,
            "output_config": outputs,
        }
    }

# Upload metadata
response = requests.post(MODELS_API_URL, json=model_metadata)
response_data = response.json()
model_uuid = response_data["data"]["model"]["uuid"]

# Upload model file
with open(model_path, "rb") as zip_file:
    files = {
        "model_file": ("model.zip", zip_file, "application/zip")
    }
    response = requests.post(url=f"{AI_TOOLKIT_BACKEND_URL}/models/{model_uuid}", # Model identifier is part of the request URL
    files=files)

    if response.status_code == 200:
        response_data = response.json()
        print(f'Pytorch model named "{model_name}" has been uploaded successfully!')
    else:
        print("Something went wrong while uploading the model: \n\n", file=sys.stderr)
        print(response.text, file=sys.stderr)

After uploading the model metadata and file, you can verify its registration and readiness status using the following endpoint. If the status remains in_progress, repeat the check until it changes to ready.

response = requests.get(f"{AI_TOOLKIT_BACKEND_URL}/models/{model_uuid}")
data = response.json()
print(f'Model status: {data["data"]["model"]["status"]}')
print(f'Model status description: {data["data"]["model"]["status_description"]}')

Conversion of PyTorch model to ONNX

Now the model is ready, and you can run the conversion.

In the code below, some parameters are set for the OnnxConversion pass. To view the full list of supported parameters, check this endpoint:

available_passes_response = requests.get(f"{AI_TOOLKIT_BACKEND_URL}/optimizations/passes")
available_passes = available_passes_response.json()

# This prints configuration parameters only for OnnxConversion pass. Feel free to change it and explore
# other passes as well.
onnx_conversion_pass_config = next(_pass for _pass in available_passes["data"]["passes"] if _pass["type"] == "OnnxConversion")
print(onnx_conversion_pass_config)

To run the conversion, send a request to the /optimizations/run endpoint:

OPTIMIZATIONS_API_URL = f"{AI_TOOLKIT_BACKEND_URL}/optimizations"
RUN_OPTIMIZATION_API_URL = f"{OPTIMIZATIONS_API_URL}/run"

pass_config = {
    "model_uuid": model_uuid,
    "passes": [
        {
            "type": "OnnxConversion",
            "config": {
                "target_opset": 14,
                "optimize": True
            }
        }
    ]
}

optimization_response = requests.post(RUN_OPTIMIZATION_API_URL, json=pass_config)
data = optimization_response.json()
optimization_uuid = data["data"]["optimization"]["uuid"]

Running the code above starts the conversion. To check the status, use the following endpoint. Again, you may need to run it multiple times until status changes to success:

response = requests.get(f"{OPTIMIZATIONS_API_URL}/{optimization_uuid}")
data = response.json()
status = data["data"]["optimization"]["status"]
print(f"Conversion status: {status}")

if status == "success":
    artifact_id = data["data"]["optimization"]["artifacts"][0]["artifact_id"]

Download converted model

# Change model path to your location
dest_model_path = Path("my_path_to_converted_model.onnx")

download_response = requests.get(f"{AI_TOOLKIT_BACKEND_URL}/optimizations/{optimization_uuid}/resources/{artifact_id}")
with dest_model_path.open("wb") as f:
    f.write(download_response.content)