Onnxruntime Test Error after Successfully Converting Model to ONNX

Question

I have a simple PyTorch model that I'm attempting to convert to ONNX format. The forward() function consists of a call to nn.transformer.encoder(). The conversion to ONNX using torch.onnx.export() completes successfully. However, when I test the model using onnxruntime_test, it fails, except for specific input cases.

I suspect the problem is related to dynamic axes and reshaping during runtime, but I'm unsure of the exact cause.

I've provided a minimal example below where I create the model, run it with two different tensors, export it to ONNX, and attempt to replicate the process using ONNX.

Any insights into why this issue might be occurring would be greatly appreciated. Thank you!

#!/usr/bin/env python3

import torch.nn as nn
from torch import Tensor, rand
import torch.onnx
import onnx

onnx_model = 'MicroTest.onnx'

class Test_trans(nn.Module):
    def __init__( self, emb_size=100):
        super(Test_trans, self).__init__()
        self.transformer = nn.Transformer(emb_size, 2, 2, 2, 512, 0.1)

    def forward(self, src: Tensor):
        return self.transformer.encoder(src)

def process_one_torch(session, ten):
    print('Tensor In size:', ten.size(), end='\t')
    memory = session(ten)
    print('Mem size:', memory.size());

def process_one_onnx(session, npa):
    ortvalue = onnxruntime.OrtValue.ortvalue_from_numpy(npa)
    print('In ortvalue.shape:', ortvalue.shape(), end='\t')
    memory = session.run(None, {session.get_inputs()[0].name: ortvalue})
    print('ONNX mem.shape:', memory[0].shape)

mini = Test_trans()

c_tensor_12 = rand((12,1,100))
c_tensor_10 = rand((10,1,100))

print('################################# Torch ###################################')
process_one_torch(mini, c_tensor_12)
process_one_torch(mini, c_tensor_10)

torch.onnx.export(mini, # model being run
    c_tensor_12,        # model input (or a tuple for multiple inputs)
    onnx_model,         # where to save the model (can be a file or file-like object)
    export_params=True, # store the trained parameter weights inside the model file
    opset_version=17,   # the ONNX version to export the model to
    do_constant_folding=False,
    input_names = ['input'],   # the model's input names
    output_names = ['output'], # the model's output names
    dynamic_axes = {'input' : {0: 'max_len'}})

print('################################# ONNX RT #################################')
import onnxruntime

session = onnxruntime.InferenceSession(onnx_model, providers=["CPUExecutionProvider"])
print('Session inputs:', session.get_inputs()[0])

process_one_onnx(session, c_tensor_12.numpy())
process_one_onnx(session, c_tensor_10.numpy()) #This one crashes
 

Running the onnxruntime_test with the created onnx model I got: ` onnxruntime_test MicroTest.onnx 2024-06-12 17:11:22.877402346 [E:onnxruntime:, sequential_executor.cc:514 ExecuteKernel] Non-zero status code returned while running Reshape node. Name:'/encoder/layers.0/self_attn/Reshape_4' Status Message: /croot/onnxruntime_1711063034809/work/onnxruntime/core/providers/cpu/tensor/reshape_helper.h:44 onnxruntime::ReshapeHelper::ReshapeHelper(const onnxruntime::TensorShape&, onnxruntime::TensorShapeVector&, bool) input_shape_size == size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{1,1,100}, requested shape:{12,2,50}

Traceback (most recent call last): File "/home/if/miniconda3/envs/cpu/bin/onnxruntime_test", line 11, in sys.exit(main()) ^^^^^^ File "/home/if/miniconda3/envs/cpu/lib/python3.11/site-packages/onnxruntime/tools/onnxruntime_test.py", line 159, in main exit_code, _, _ = run_model(args.model_path, args.num_iters, args.debug, args.profile, args.symbolic_dims) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/if/miniconda3/envs/cpu/lib/python3.11/site-packages/onnxruntime/tools/onnxruntime_test.py", line 118, in run_model outputs = sess.run([], feeds) # fetch all outputs ^^^^^^^^^^^^^^^^^^^ File "/home/if/miniconda3/envs/cpu/lib/python3.11/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 220, in run return self._sess.run(output_names, input_feed, run_options) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ onnxruntime.capi.onnxruntime_pybind11_state.RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running Reshape node. Name:'/encoder/layers.0/self_attn/Reshape_4' Status Message: /croot/onnxruntime_1711063034809/work/onnxruntime/core/providers/cpu/tensor/reshape_helper.h:44 onnxruntime::ReshapeHelper::ReshapeHelper(const onnxruntime::TensorShape&, onnxruntime::TensorShapeVector&, bool) input_shape_size == size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{1,1,100}, requested shape:{12,2,50} ` All I understand is that inside is reshaped something differently then expected.

Answer 1

O.K. This is clearly a version issue with the ONNX exporter. I tried an older version of pytorch.onnx.export(), and everything runs smoothly, with onnxruntime_test producing no errors. Unfortunately, the older version only supports opset 15, which leads to other problems, but the main issue is resolved.