Add in support for 1 unknown dimension for ONNX runtime. (#6265)

Author: michaelgsharp

src/Microsoft.ML.OnnxTransformer/OnnxTransform.cs

@@ -385,8 +385,13 @@ private static IEnumerable<int> AdjustDimensions(OnnxShape shape)
         {
             if (shape.Count > 0)
             {
-                return shape.Select(x => (x <= 0) ? 1 : x);
+                if (shape[0] < 0)
+                {
+                    shape[0] = 1;
+                }
+                return shape.Select(x => (x <= 0) ? 0 : x);
             }
+
             return new[] { 1 };
         }
 
@@ -444,6 +449,11 @@ public Mapper(OnnxTransformer parent, DataViewSchema inputSchema) :
                     var shape = inputNodeInfo.Shape;
 
                     var inputShape = AdjustDimensions(inputNodeInfo.Shape);
+
+                    // Only allow a single unkown size dimension
+                    if (inputShape.Where(x => x == 0).Count() > 1)
+                        throw new ArgumentOutOfRangeException(_parent.Inputs[i], "Only 1 unknown dimension is allowed");
+
                     _inputTensorShapes[i] = inputShape.ToList();
                     _inputOnnxTypes[i] = inputNodeInfo.TypeInOnnxRuntime;
 
@@ -456,9 +466,6 @@ public Mapper(OnnxTransformer parent, DataViewSchema inputSchema) :
                     var type = inputSchema[_inputColIndices[i]].Type;
                     var vectorType = type as VectorDataViewType;
 
-                    if (vectorType != null && vectorType.Size == 0)
-                        throw Host.Except($"Variable length input columns not supported");
-
                     var itemType = type.GetItemType();
                     var nodeItemType = inputNodeInfo.DataViewType.GetItemType();
                     if (itemType != nodeItemType)
@@ -474,11 +481,14 @@ public Mapper(OnnxTransformer parent, DataViewSchema inputSchema) :
 
                     // If the column is one dimension we make sure that the total size of the Onnx shape matches.
                     // Compute the total size of the known dimensions of the shape.
-                    int valCount = inputShape.Where(x => x > 0).Aggregate((x, y) => x * y);
-                    // The column length should be divisible by this, so that the other dimensions can be integral.
-                    int typeValueCount = type.GetValueCount();
-                    if (typeValueCount % valCount != 0)
-                        throw Contracts.Except($"Input shape mismatch: Input '{_parent.Inputs[i]}' has shape {String.Join(",", inputShape)}, but input data is of length {typeValueCount}.");
+                    if (!inputShape.Any(x => x == 0))
+                    {
+                        int valCount = inputShape.Where(x => x > 0).Aggregate((x, y) => x * y);
+                        // The column length should be divisible by this, so that the other dimensions can be integral.
+                        int typeValueCount = type.GetValueCount();
+                        if (typeValueCount % valCount != 0)
+                            throw Contracts.Except($"Input shape mismatch: Input '{_parent.Inputs[i]}' has shape {String.Join(",", inputShape)}, but input data is of length {typeValueCount}.");
+                    }
                 }
             }
 
@@ -781,23 +791,56 @@ public NamedOnnxValue GetNamedOnnxValue()
 
             private class NamedOnnxValueGetterVec<T> : INamedOnnxValueGetter
             {
+                private delegate NamedOnnxValue GetNamedOnnxVal();
+
                 private readonly ValueGetter<VBuffer<T>> _srcGetter;
                 private readonly OnnxShape _tensorShape;
                 private readonly string _colName;
                 private VBuffer<T> _vBuffer;
                 private VBuffer<T> _vBufferDense;
+                private readonly int _denominator;
+                private readonly int _zeroIndex;
+                private readonly GetNamedOnnxVal _namedOnnxValueDelegate;
+
                 public NamedOnnxValueGetterVec(DataViewRow input, int colIndex, OnnxShape tensorShape)
                 {
                     _srcGetter = input.GetGetter<VBuffer<T>>(input.Schema[colIndex]);
                     _tensorShape = tensorShape;
                     _colName = input.Schema[colIndex].Name;
                     _vBuffer = default;
                     _vBufferDense = default;
+                    _denominator = _tensorShape.Where(x => x > 0).Aggregate((a, x) => a * x);
+                    _zeroIndex = _tensorShape.IndexOf(0);
+
+                    var isKnownSize = (input.Schema[colIndex].Type as VectorDataViewType).IsKnownSize;
+
+                    if (isKnownSize)
+                        _namedOnnxValueDelegate = GetNamedOnnxValueKnownSize;
+                    else
+                        _namedOnnxValueDelegate = GetNamedOnnxValueUnknownSize;
                 }
                 public NamedOnnxValue GetNamedOnnxValue()
+                {
+                    return _namedOnnxValueDelegate();
+                }
+
+                private void GetNamedOnnxValueCore()
                 {
                     _srcGetter(ref _vBuffer);
                     _vBuffer.CopyToDense(ref _vBufferDense);
+                }
+
+                private NamedOnnxValue GetNamedOnnxValueKnownSize()
+                {
+                    GetNamedOnnxValueCore();
+                    return OnnxUtils.CreateNamedOnnxValue(_colName, _vBufferDense.GetValues(), _tensorShape);
+                }
+
+                private NamedOnnxValue GetNamedOnnxValueUnknownSize()
+                {
+                    GetNamedOnnxValueCore();
+
+                    _tensorShape[_zeroIndex] = _vBufferDense.Length / _denominator;
                     return OnnxUtils.CreateNamedOnnxValue(_colName, _vBufferDense.GetValues(), _tensorShape);
                 }
             }
@@ -908,14 +951,14 @@ public override SchemaShape GetOutputSchema(SchemaShape inputSchema)
                 // Get the i-th IDataView input column's name in the underlying ONNX transformer.
                 var input = Transformer.Inputs[i];
 
+                // Only allow 1 unknown dimension
+                if (Transformer.Model.ModelInfo.InputsInfo[i].Shape.Where(x => x == 0).Count() > 1)
+                    throw new ArgumentOutOfRangeException(input, "Only 1 unknown dimension is allowed");
+
                 // Make sure inputSchema contains the i-th input column.
                 if (!inputSchema.TryFindColumn(input, out var col))
                     throw Host.ExceptSchemaMismatch(nameof(inputSchema), "input", input);
 
-                // Make sure that the input columns in inputSchema are fixed shape tensors.
-                if (col.Kind == SchemaShape.Column.VectorKind.VariableVector)
-                    throw Host.ExceptSchemaMismatch(nameof(inputSchema), "input", input, "vector", col.GetTypeString());
-
                 var inputsInfo = Transformer.Model.ModelInfo.InputsInfo;
                 var idx = Transformer.Model.ModelInfo.InputNames.IndexOf(input);
                 if (idx < 0)