GitHub Repository: ibm/watson-machine-learning-samples
Path: blob/master/cloud/notebooks/python_sdk/converters/Convert ONNX neural network from fixed axes to dynamic axes.ipynb
⁶⁴⁰⁵ views

Kernel: Python 3 (ipykernel)

Convert ONNX neural network from fixed axes to dynamic axes and use it with `ibm-watsonx-ai`

This notebook facilitates ONNX, PyTorch and watsonx.ai Runtime service. It contains steps and code to work with ibm-watsonx-ai library available in PyPI repository. It introduces commands for adapting model to dynamic axes, getting model and training data, persisting model, deploying model and scoring it.

Some familiarity with Python is helpful. This notebook uses Python 3.11.

Learning goals

The learning goals of this notebook are:

Convert neural network from fixed axes to dynamic axes.
Persist converted model in watsonx.ai Runtime repository.
Deploy model for online scoring using client library.
Score sample records using client library.

This notebook contains the following parts:

1. Setting up the environment

Before you use the sample code in this notebook, you must perform the following setup tasks:

Create a watsonx.ai Runtime instance (information on service plans and further reading can be found here).

1.1. Installing and importing the `ibm-watsonx-ai` and dependecies

Note: ibm-watsonx-ai documentation can be found here.

In [ ]:

!pip install -U ibm-watsonx-ai | tail -n 1
!pip install torch==2.1 | tail -n 1
!pip install onnx==1.16 | tail -n 1
!pip install onnxruntime==1.16.3 | tail -n 1

In [2]:

import getpass
import json
import torch
import torch.nn as nn
import onnx
import onnxruntime as ort
from onnx.tools import update_model_dims
from ibm_watsonx_ai import Credentials, APIClient

1.2. Connecting to watsonx.ai Runtime

Authenticate to the watsonx.ai Runtime service on IBM Cloud. You need to provide platform api_key and instance location.

You can use IBM Cloud CLI to retrieve platform API Key and instance location.

API Key can be generated in the following way:

ibmcloud login
ibmcloud iam api-key-create API_KEY_NAME

In result, get the value of api_key from the output.

Location of your watsonx.ai Runtime instance can be retrieved in the following way:

ibmcloud login --apikey API_KEY -a https://cloud.ibm.com
ibmcloud resource service-instance INSTANCE_NAME

In result, get the value of location from the output.

Tip: Your Cloud API key can be generated by going to the Users section of the Cloud console. From that page, click your name, scroll down to the API Keys section, and click Create an IBM Cloud API key. Give your key a name and click Create, then copy the created key and paste it below. You can also get a service specific url by going to the Endpoint URLs section of the watsonx.ai Runtime docs. You can check your instance location in your watsonx.ai Runtime Service instance details.

You can also get service specific apikey by going to the Service IDs section of the Cloud Console. From that page, click Create, then copy the created key and paste it below.

Action: Enter your api_key and location in the following cells.

In [2]:

api_key = getpass.getpass("Please enter your api key (hit enter): ")

Out[2]:

Please enter your api key (hit enter):  ········

In [2]:

location = 'ENTER YOUR LOCATION HERE'

If you are running this notebook on Cloud, you can access the location via:

location = os.environ.get("RUNTIME_ENV_REGION")

In [3]:

credentials = Credentials(
    api_key=api_key,
    url=f'https://{location}.ml.cloud.ibm.com'
)

In [4]:

client = APIClient(credentials=credentials)

1.3. Working with spaces

First of all, you need to create a space that will be used for your work. If you do not have space already created, you can use Deployment Spaces Dashboard to create one.

Click New Deployment Space
Create an empty space
Select Cloud Object Storage
Select watsonx.ai Runtime instance and press Create
Copy space_id and paste it below

Tip: You can also use the ibm_watsonx_ai SDK to prepare the space for your work. More information can be found here.

Action: Assign space ID below

In [5]:

space_id = 'ENTER YOUR SPACE ID HERE'

You can use list method to print all existing spaces.

In [ ]:

client.spaces.list(limit=10)

To be able to interact with all resources available in watsonx.ai Runtime, you need to set space which you will be using.

In [21]:

client.set.default_space(space_id)

Out[21]:

'SUCCESS'

2. (OPTIONAL) Creating and exporting basic ONNX model

This optional section demonstrates exporting a simple PyTorch model to the ONNX format. The model is configured with fixed tensor axes (i.e., dynamic_axes are not defined), meaning it only supports a batch size of 1. The example serves as a placeholder to illustrate the process of enabling dynamic tensor axes in ONNX models later. You may replace this example model with your own or download a model that meets your specific needs.

In [3]:

onnx_model_path = "model.onnx"

In [4]:

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 1)

    def forward(self, x):
        return self.linear(x)

model = SimpleModel()

torch_input = torch.randn(1, 2) 

torch.onnx.export(model,
                  torch_input,
                  onnx_model_path,
                  verbose=False,
                  export_params=True,
                  keep_initializers_as_inputs=True,
                  input_names=['input'],
                  output_names=['output'])

print(f"ONNX model has been saved to {onnx_model_path}")

Out[4]:

ONNX model has been saved to model.onnx

Initialize an ONNX Runtime session and run inference on a single batch

In [5]:

ort.set_default_logger_severity(3)

ort_session = ort.InferenceSession(onnx_model_path, providers=['CPUExecutionProvider'])

print("Original Input Shapes:")
for i in ort_session.get_inputs():
    print(f"Input Name: {i.name}, Shape: {i.shape}")

print("\nOriginal Output Shapes:")
for o in ort_session.get_outputs():
    print(f"Output Name: {o.name}, Shape: {o.shape}")

input_data = {ort_session.get_inputs()[0].name: torch_input.numpy()}
print(f"\nPrediciton result: {ort_session.run([], input_data)}")

Out[5]:

Original Input Shapes:
Input Name: input, Shape: [1, 2]

Original Output Shapes:
Output Name: output, Shape: [1, 1]

Prediciton result: [array([[-0.848896]], dtype=float32)]

Attempt multi-batch inference, expected to fail due to fixed axes

In [6]:

batch_torch_input = torch.randn(3, 2) 
batch_input_data = {ort_session.get_inputs()[0].name: batch_torch_input.numpy()}

try:
    batch_output = ort_session.run([], batch_input_data)
    print("Multi-batch inference result:", batch_output)
except Exception as e:
    print("Multi-batch inference failed as expected:", str(e))

Out[6]:

Multi-batch inference failed as expected: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Got invalid dimensions for input: input for the following indices
 index: 0 Got: 3 Expected: 1
 Please fix either the inputs or the model.

3. Converting model from fixed axes to dynamic axes

In this section, you will convert an ONNX neural network with fixed axes to a model that supports dynamic axes. This allows for flexible input and output dimensions, making it possible to run inference with varying batch sizes.

To use your own model, provide its file path in onnx_custom_model_path. Otherwise, the default path from the second section will be used.

In [7]:

onnx_custom_model_path = ''

In [8]:

if onnx_custom_model_path:
    onnx_model_path = onnx_custom_model_path

3.1. Updating Model to support dynamic batch sizes

You will now update the model to support dynamic axes for the first dimension (batch size) on all inputs and outputs.

In [9]:

model = onnx.load(onnx_model_path)

Using model.graph here is necessary because update_inputs_outputs_dims operates on this structure. Note that model.graph may contain additional information beyond the primary inputs and outputs shape. Therefore, dynamic conversion is performed on the output from InferenceSession, which accurately reflects the exact inputs and outputs.

Initialize dictionaries to hold input and output dimensions from the model's graph

In [10]:

input_dims = {i.name:  [d.dim_value for d in i.type.tensor_type.shape.dim] for i in model.graph.input}
output_dims = {o.name:  [d.dim_value for d in o.type.tensor_type.shape.dim] for o in model.graph.output}

Loop over all inputs/outputs and update the first axis to dynamic while keeping other axes the same

In [11]:

for inpt in ort_session.get_inputs():
    input_dims[inpt.name] = ["batch_size"] + inpt.shape[1:]

for output in ort_session.get_outputs():
    output_dims[output.name] = ["batch_size"] + output.shape[1:]

updated_model = update_model_dims.update_inputs_outputs_dims(model, input_dims, output_dims)

dynamic_model_path = "dynamic_model.onnx"
onnx.save(updated_model, dynamic_model_path)

Initialize an updated ONNX Runtime session

In [12]:

dynamic_ort_session = ort.InferenceSession(dynamic_model_path, providers=['CPUExecutionProvider'])

print("Original Input Shapes:")
for i in dynamic_ort_session.get_inputs():
    print(f"Input Name: {i.name}, Shape: {i.shape}")

print("\nOriginal Output Shapes:")
for o in dynamic_ort_session.get_outputs():
    print(f"Output Name: {o.name}, Shape: {o.shape}")

Out[12]:

Original Input Shapes:
Input Name: input, Shape: ['batch_size', 2]

Original Output Shapes:
Output Name: output, Shape: ['batch_size', 1]

Run inference on the updated model with a batch input

3.2. Evaluating the ONNX Model

In [16]:

dynamic_ort_session.run([], batch_input_data)

Out[16]:

[array([[-0.10204538],
        [-0.63963014],
        [ 0.06598502]], dtype=float32)]

The model now supports variable batch sizes, allowing for flexible batch scoring during inference.

4. Persisting adjusted ONNX model

In this section, you will learn how to store your converted ONNX model in watsonx.ai Runtime repository using the IBM watsonx.ai SDK.

4.1. Publishing model in watsonx.ai Runtime repository

Define model name, type and software spec.

In [17]:

sofware_spec_id = client.software_specifications.get_id_by_name("onnxruntime_opset_19")
onnx_model_zip = "onnx_dynamic_model.zip"

In [18]:

!zip {onnx_model_zip} {dynamic_model_path}

Out[18]:

updating: dynamic_model.onnx (deflated 31%)

In [22]:

metadata = {
    client.repository.ModelMetaNames.NAME: 'Dynamic axes ONNX model',
    client.repository.ModelMetaNames.TYPE: 'onnxruntime_1.16',
    client.repository.ModelMetaNames.SOFTWARE_SPEC_ID: sofware_spec_id
}

published_model = client.repository.store_model(
    model=onnx_model_zip,
    meta_props=metadata
)

4.2. Getting model details

In [ ]:

published_model_id = client.repository.get_model_id(published_model)
model_details = client.repository.get_details(published_model_id)
print(json.dumps(model_details, indent=2))

5. Deploying and scoring ONNX model

In this section you'll learn how to create an online scoring service and predict on unseen data.

5.1. Creating online deployment for published model

In [24]:

metadata = {
    client.deployments.ConfigurationMetaNames.NAME: "Deployment of dynamic axes ONNX model",
    client.deployments.ConfigurationMetaNames.ONLINE: {}
}

created_deployment = client.deployments.create(published_model_id, meta_props=metadata)

Out[24]:

######################################################################################

Synchronous deployment creation for id: 'dd3a020b-315d-43bd-932e-19c3c6584dcc' started

######################################################################################


initializing
Note: online_url and serving_urls are deprecated and will be removed in a future release. Use inference instead.
..
ready


-----------------------------------------------------------------------------------------------
Successfully finished deployment creation, deployment_id='3c9f62d4-5344-4cc7-9212-783395945540'
-----------------------------------------------------------------------------------------------

In [25]:

deployment_id = client.deployments.get_id(created_deployment)

Now you can print an online scoring endpoint.

In [ ]:

client.deployments.get_scoring_href(created_deployment)

5.2. Getting deployment details

In [ ]:

client.deployments.get_details(deployment_id)

5.3. Scoring

You can use below method to do test scoring request against deployed model.

In [29]:

scoring_payload = {"input_data": [{"values": batch_torch_input.tolist()}]}

Use client.deployments.score() method to run scoring.

In [30]:

predictions = client.deployments.score(deployment_id, scoring_payload)

Let's print the result of predictions.

In [31]:

print(json.dumps(predictions, indent=2))

Out[31]:

{
  "predictions": [
    {
      "id": "output",
      "values": [
        [
          -0.10204537957906723
        ],
        [
          -0.639630138874054
        ],
        [
          0.06598502397537231
        ]
      ]
    }
  ]
}

As you can see, predicted values are the same one as received in the evaluation part from test dataset.

6. Cleaning up

If you want to clean up after the notebook execution, i.e. remove any created assets like:

experiments
trainings
pipelines
model definitions
models
functions
deployments

please follow up this sample notebook.

7. Summary and next steps

You successfully completed this notebook! You learned how to use ONNX library as well as watsonx.ai Runtime for model conversion and deployment. Check out our Online Documentation for more samples, tutorials, documentation, how-tos, and blog posts.

Authors

Michał Koruszowic, Software Engineer

Convert ONNX neural network from fixed axes to dynamic axes and use it with `ibm-watsonx-ai`

Learning goals

Contents

1. Setting up the environment

1.1. Installing and importing the `ibm-watsonx-ai` and dependecies

1.2. Connecting to watsonx.ai Runtime

1.3. Working with spaces

2. (OPTIONAL) Creating and exporting basic ONNX model

3. Converting model from fixed axes to dynamic axes

3.1. Updating Model to support dynamic batch sizes

3.2. Evaluating the ONNX Model

4. Persisting adjusted ONNX model

4.1. Publishing model in watsonx.ai Runtime repository

4.2. Getting model details

5. Deploying and scoring ONNX model

5.1. Creating online deployment for published model

5.2. Getting deployment details

5.3. Scoring

6. Cleaning up

7. Summary and next steps

Authors

Product

Resources

Company

Convert ONNX neural network from fixed axes to dynamic axes and use it with ibm-watsonx-ai

Learning goals

Contents

1. Setting up the environment

1.1. Installing and importing the ibm-watsonx-ai and dependecies

1.2. Connecting to watsonx.ai Runtime

1.3. Working with spaces

2. (OPTIONAL) Creating and exporting basic ONNX model

3. Converting model from fixed axes to dynamic axes

3.1. Updating Model to support dynamic batch sizes

3.2. Evaluating the ONNX Model

4. Persisting adjusted ONNX model

4.1. Publishing model in watsonx.ai Runtime repository

4.2. Getting model details

5. Deploying and scoring ONNX model

5.1. Creating online deployment for published model

5.2. Getting deployment details

5.3. Scoring

6. Cleaning up

7. Summary and next steps

Authors

Convert ONNX neural network from fixed axes to dynamic axes and use it with `ibm-watsonx-ai`

1.1. Installing and importing the `ibm-watsonx-ai` and dependecies