Opaque disruptive handling of variable axes of ONNX-models

[UlrichSuess]

### System information

• Win 10, 1903
• .NET 4.7.2

### Issue
remark: I do have a functional workaround solution, but it requires quite some nasty additional code and probably is a bit slower than optimal.

The main point of critique here concerns the OnnxTransform class.

What did you do?
Tried to do inference with an ONNX-model for sequence labeling with two variable input axes (batch and sequence).

The whole input (3 dimensions: batch, sequence, features) were put into one big float[]. Still not sure whether this is really the correct way. All attempts to use multidimensional input fields failed.

What happened?
As stated in source code (OnnxTransform.MakeRowMapper(...), see below), dimensions of variable input axes were set to 1, resulting in incorrect batches and sequences of size 1 each, regardless of specified DataView.

There is no explicit hint leading to the dimensionality reduction, however. The problem is documented, but that's in a private method. You only get an exception when feeding your whole sequence (as stated as one big float[]), because due to the dimension reduction the set up model expects length 1 * 1 * FeatureLength, and the input vector gets is large for that. At least that exception tells you that the product of the dimensions is relevant.

The workaround consists of setting up the "model" for each sequence by using the ApplyOnnxModel-overload with the input size dictionary "shapesDictionary". The documentation for that overload method is a bit cryptic, though (see below). It does allow to override the reduction to 1 of the variable axes (and only for those!) by using the values in the dictionary, but I learned that from the source code, not from the documentation.

The DataView has to be reset for each sequence, too, of course, to contain the specific sequence length.

Took unnecessarily long to find the source of the problem and the workaround solution.

What did you expect?
Best case: Full support for variable input dimensions, which should extend to the DataView and the model.

2nd best: Even for my chosen workaround of adjusting everything for each sequence to contain specific sequence length, the model-row-mapper should infer the correct dimensions instead of reducing them to 1 (there IS a todo comment for that). Fortunately, the time for the resets still make the workaround kind of usable. (although at batch size 1, so far, for more padding will probably be necessary, the effect of which remains to be checked)

At least: Output some kind of hint that the reduction takes place, at least in the public documentation for ApplyOnnxModel(...) or with further information in the exception when the sizes don't fit.

Source code / logs

From OnnxTransform.cs:

        private protected override IRowMapper MakeRowMapper(DataViewSchema inputSchema) => new Mapper(this, inputSchema);

        /// <summary>
        /// This design assumes that all unknown dimensions are 1s. It also convert scalar shape [] in ONNX to [1].
        /// [TODO] We should infer the unknown shape from input data instead of forcing them to be 1.
        /// </summary>
        private static IEnumerable<int> AdjustDimensions(OnnxShape shape)
        {
            if (shape.Count > 0)
            {
                return shape.Select(x => (x <= 0) ? 1 : x);
            }
            return new[] { 1 };
        }

From OnnxCatalog.ApplyOnnxModel(...) not too helpful documentation:

        /// <param name="shapeDictionary">ONNX shape should be used to over those loaded from <paramref name="modelFile"/>.</param>

suggestion:

        /// <param name="shapeDictionary">Variable dimension sizes as stated in the loaded <paramref name="modelFile"/> are replaced by dimension sizes from this dictionary. All dimensions, also nonvariable, have to be given, though. For keys use names as stated in the ONNX model, e.g. "input"</param>

[antoniovs1029]

Hi, @UlrichSuess . Thank you very much for your feedback. I'm very interested in the workaround you mention, and have some questions I'll make below. But first, let me explain some points:

1. OnnxTransformer doesn't really support running batch inferencing (i.e., it always runs with a batch size of 1). The OnnxTransformer in ML.NET is pretty much a wrapper around OnnxRuntime, and although I think OnnxRuntime supports batch inferencing, OnnxTransformer currently doesn't. The reason behind this is that ML.NET runs inferencing on one sample (i.e. "one row" of an input IDataView) at a time, and is designed and optimized to do this... so the concept of "batch" doesn't really exist in ML.NET. Inside the team, we've discussed enabling batch inferencing for the OnnxTransformer, but some work is required to fit "batches" into the architecture used by ML.NET... and we don't have any plan to add this functionality soon. Also, this isn't a problem for the general user, because we automatically set the batch dimension (typically the "-1" or "None" found at the beginning of a shape) to 1 (this is done in the method you've pointed to: AdjustDimension(), but also in other methods like GetTensorDims()).

2. OnnxTransformer doesn't infer -1/ None dimensions and assumes they're 1 by default. As you've correctly pointed, this pretty much happens in the AdjustDimension() method, and the workaround for users is to use a shapeDictionary (introduced in Allow user to overwrite unknown shapes loaded from ONNX model #3963 precisely for this use case). Currently there's no plan to implement automatic inference of this kind of dimensions, but we'll take your suggestions into account when planning future changes.

3. ML.NET only supports unidimensional arrays as inputs, although they can have multidimensional VectorTypeAttributes. So this is invalid in ML.NET:

[VectorType(2, 3)]
public float[,] input { get; set; } // use arrays of dimensions 2, 3

And actually what ML.NET needs is:

[VectorType(2, 3)]
public float[] input { get; set; } // use arrays of size 6

This means ML.NET does expect the user to "flatten" their input multidimensional array into a one-dimensional array.

4. About ML.NET docs and exception messages. So I identify two main areas of improvement in our docs / exception messages:
   4.1 ApplyOnnxModel Clearing up the usage of shapeDictionary seems necessary in our docs. I'll try to get the suggestion you've made about this into the docs for the next release of ML.NET.
   4.2 VectorTypeAttribute I think the exception of using a multidimensional array as input of ML.NET is System.NullReferenceException: 'Object reference not set to an instance of an object.' and I think we can improve this exception message to be more clear that multidimensional arrays aren't allowed in ML.NET (see Error on prediction with LSTM Model in ONNX Format in ML.Net #5273 (comment)) . Also I'm considering in modifying some docs regarding this.

Regarding your workaround.

So I was wondering if it's possible for you to share a .zip file containing your code, onnx model and some sample data to understand the workaround you're describing. This could be helpful to understand more about how users use Onnx in ML.NET, and also provide us ideas when we get to plan future enhancements to the OnnxTransformer.

If you're not able to share it, then a description of the shapes of the input and output of your onnx model, along with the code of your input and output classes and (perhaps) the code of how you flatten your input array, would be useful.

In particular I'd be interested in these points:
1. Does your workaround let you run multiple batches at once?. As I've explained above I think this would be impossible as OnnxTransformer uses a batch size = 1, but I might be wrong about this being impossible, and maybe your workaround accomplishes it. Reading your issue it seems your workaround also sets the batch dimension to 1, but I just want to confirm.

2. Does your workaround let you use variable length "sequences"?. From reading your issue I'm not sure if this is required to you, and if your workaround handles this. In particular, what do you mean when you say "The DataView has to be reset for each sequence, too, of course, to contain the specific sequence length." Does every row in your input dataset have "sequences" of different size?

3. Is your output of variable length too? If so, how do you handle it?

[UlrichSuess]

Hi Antonio,

thanks for your fast and very informative response.

Alas, I cannot give you the project itself, there is company IP in the model and test data, but probably that won't be necessary, anyway. The code is in a rather unpresentable proof of concept state, too, right now.

For starters, though, here a few explanations and code snippets hopefully clarifying the approach.

the code of how you flatten your input array, would be useful.

After unsuccessfully trying multidimensional approaches, in the end we arrived at the simple but sufficient input class

    public class InputSequence
    {
        // the values need to be constants and are just named to clarify their purpose, 
        // BUT FOR THE WORKAROUND THE VECTOR TYPE IS RESET LATER, ANYWAY
        [VectorType(batchSize, sequenceLength, featuresLength)]
        public float[] input;
    }

There is nothing special about the flattening itself. There is a data loader class that produces InputSequence objects from a test set file, right now always assuming batch size 1 and just concatenating the features of all sequence tokens into one array.

This is stuffed into an IEnumerable called "sequence" and used later on (see following section).

2. Does your workaround let you use variable length "sequences"?

As for variable sequences and setting up DataView and the model: for the batchSize=1-scenario consider the following snippets:

for the DataView:

int sequLength = sequence.First().input.Length / featuresLength;
var schemaDef = SchemaDefinition.Create(typeof(InputSequence));
schemaDef["input"].ColumnType = new VectorDataViewType(NumberDataViewType.Single, 1, sequLength, featuresLength);
// sequence is an IEnumerable<InputSequence> containing the actual sequence data in an flattened array
var dataViewSequence = context.Data.LoadFromEnumerable(sequence, schemaDef);

for the model:

shapesDictionary = new Dictionary<string, int[]>
{
    {"input", new int[] { batchSize, sequenceLength, feauturesLength }
};

IEstimator<ITransformer> modelTransformer = mlContext.Transforms.ApplyOnnxModel(
        modelFile: modelLocation,
        outputColumnNames: new[] {ModelConfig.ModelData.ModelOutputName},
        inputColumnNames: new[] {ModelConfig.ModelData.ModelInputName},
        shapeDictionary: shapesDictionary);

This extra code has to be run for each sequence before actually inferring from the model. So there is no real variable length, right now it is rather "always reset to the proper fixed size".

There are loads of task cancel exceptions, by the way, when running this through a few hundred sequences. These exceptions are caught by ML.NET-code, apparently, and cause no visible functional problems.

Does every row in your input dataset have "sequences" of different size?

Yes, our input data has sequences of different size, but as stated above, the DataView is reset for each new sequence (in the batchSize=1-scenario).

1. Does your workaround let you run multiple batches at once?

I have just been able to extend the workaround to use batch sizes > 1 with a very useful speedup in raw inference time at a currently rather high cost for additional data preparation, especially padding shorter sequenes. In the specific case, batch size 10 reduces total inference time to a third. Work in progress. By the way, all on cpu. Using batch size > 1 in this case means setting up the model and the DataView only once per whole batch, with all sequences padded to the length of the longest in the batch. So for a batch, all sequences have to be of same length. That is the same as in PyTorch, though.

3. Is your output of variable length too?

Our output contains one large float[]. Knowing the output size (here four labels), each sequence length (unpadded!) within the current batch and the current batch size (which may vary, even for intended fixed size, as the very last badge may be smaller), it's easy to split that to obtain the proper outputs for each sequence.
I have not tried to adjust dimensions for the output with the shape dictionary, yet. That might remove some splitting code.

Feel free to ask for further explanations if necessary.

Uli

EDIT:

in snippet for DataView-adjusment

• added line showing how sequence length is computed for an encountered sequence

• int the columnType-definition line: added one more int for the batchSize, here 1. This is actually not necessary as long as batch size is actually one, as internally the numbers are just multiplied to determine overall size. The model does need the correct number for the batch in the shapes dictionary, though, and for batch sizes bigger than one, the extra int becomes necessary.

[antoniovs1029]

Thanks for your detailed answer, @UlrichSuess .

Your workaround for making OnnxTransformer to work with batches is interesting, and as long as OnnxTransformer doesn't support an option that facilitates running with batches, it does sound that your workaround would be the best way to accomplish it (although it does require the preprocessing steps of requiring the user to concatenate and flatten the input arrays, as well as having to keep track of the correct shapes and VectorTypeAttribute() dimensions). It's also interesting that you report that using the workaround makes your code to run faster.

Your workaround for working with variable length dimensions is also interesting. As of now I can't think of any other way of accomplishing this on ML.NET, so maybe your approach of needing to "reset" the schema definition is the best workaround there is.

[UlrichSuess]

Thanks for your assessment, Antonio.

We will stick to this approach for the time being, then, and keep an eye on coming ML-NET changes.