Fixed OnnxTransformer output column mapping.

src/Microsoft.ML.OnnxTransformer/OnnxTransform.cs

@@ -340,6 +340,17 @@ private static IEnumerable<int> AdjustDimensions(OnnxShape shape)
             return new[] { 1 };
         }
 
+        /// <summary>
+        /// In the case that the ML.Net user wants a subset of columns or lists the columns in a different order then specified in the ONNX model,
+        /// we need to map from the ML.Net dataview column index to the ONNX model output index. This method does that mapping.
+        /// </summary>
+        /// <param name="iinfo">The index of the ML.Net column requested.</param>
+        /// <returns>The index of ONNX output.</returns>
+        internal int MapDataViewColumnToOnnxOutputTensor(int iinfo)
+        {
+            return Model.ModelInfo.OutputNames.IndexOf(Outputs[iinfo]);
+        }
+
         private sealed class Mapper : MapperBase
         {
             private readonly OnnxTransformer _parent;
@@ -468,7 +479,7 @@ protected override Delegate MakeGetter(DataViewRow input, int iinfo, Func<int, b
 
                 var activeOutputColNames = _parent.Outputs.Where((x, i) => activeOutput(i)).ToArray();
 
-                if (_parent.Model.ModelInfo.OutputsInfo[iinfo].DataViewType is VectorDataViewType vectorType)
+                if (_parent.Model.ModelInfo.OutputsInfo[_parent.MapDataViewColumnToOnnxOutputTensor(iinfo)].DataViewType is VectorDataViewType vectorType)
                 {
                     var elemRawType = vectorType.ItemType.RawType;
                     var srcNamedValueGetters = GetNamedOnnxValueGetters(input, _inputColIndices, _inputOnnxTypes, _inputTensorShapes);
@@ -479,7 +490,7 @@ protected override Delegate MakeGetter(DataViewRow input, int iinfo, Func<int, b
                 }
                 else
                 {
-                    var type = _parent.Model.ModelInfo.OutputsInfo[iinfo].DataViewType.RawType;
+                    var type = _parent.Model.ModelInfo.OutputsInfo[_parent.MapDataViewColumnToOnnxOutputTensor(iinfo)].DataViewType.RawType;
                     var srcNamedValueGetters = GetNamedOnnxValueGetters(input, _inputColIndices, _inputOnnxTypes, _inputTensorShapes);
                     return Utils.MarshalInvoke(MakeObjectGetter<int>, type, input, iinfo, srcNamedValueGetters, activeOutputColNames);
                 }
@@ -567,7 +578,7 @@ private Delegate MakeObjectGetter<T>(DataViewRow input, int iinfo, INamedOnnxVal
                     UpdateCacheIfNeeded(input.Position, srcNamedValueGetters, activeOutputColNames, outputCache);
                     var namedOnnxValue = outputCache.Outputs[_parent.Outputs[iinfo]];
                     var trueValue = namedOnnxValue.AsEnumerable<NamedOnnxValue>().Select(value => value.AsDictionary<string, float>());
-                    var caster = _parent.Model.ModelInfo.OutputsInfo[iinfo].Caster;
+                    var caster = _parent.Model.ModelInfo.OutputsInfo[_parent.MapDataViewColumnToOnnxOutputTensor(iinfo)].Caster;
                     dst = (T)caster(namedOnnxValue);
                 };
                 return valueGetter;

test/Microsoft.ML.OnnxTransformerTest/OnnxTransformTests.cs

@@ -332,6 +332,60 @@ public void OnnxModelMultiInput()
             }
         }
 
+        [OnnxFact]
+        public void OnnxModelOutputDifferentOrder()
+        {
+            var modelFile = Path.Combine(Directory.GetCurrentDirectory(), "twoinput", "twoinput.onnx");
+
+            var dataView = ML.Data.LoadFromEnumerable(
+                new TestDataMulti[] {
+                    new TestDataMulti()
+                    {
+                        ina = new float[] {1,2,3,4,5},
+                        inb = new float[] {1,2,3,4,5}
+                    }
+                });
+            // The model returns the output columns in the order outa, outb. We are doing the opposite here, making sure the name mapping is correct.
+            var onnx = ML.Transforms.ApplyOnnxModel(new[] { "outb", "outa" }, new[] { "ina", "inb" }, modelFile).Fit(dataView).Transform(dataView);
+
+            var outaCol = onnx.Schema["outa"];
+            var outbCol = onnx.Schema["outb"];
+            using (var curs = onnx.GetRowCursor(outaCol, onnx.Schema["outb"]))
+            {
+                var getScoresa = curs.GetGetter<VBuffer<float>>(outaCol);
+                var getScoresb = curs.GetGetter<VBuffer<float>>(outbCol);
+                var buffera = default(VBuffer<float>);
+                var bufferb = default(VBuffer<float>);
+
+                while (curs.MoveNext())
+                {
+                    getScoresa(ref buffera);
+                    getScoresb(ref bufferb);
+                    Assert.Equal(5, buffera.Length);
+                    Assert.Equal(5, bufferb.Length);
+                    Assert.Equal(0, buffera.GetValues().ToArray().Sum());
+                    Assert.Equal(30, bufferb.GetValues().ToArray().Sum());
+                }
+            }
+
+            // The model returns the output columns in the order outa, outb. We are doing only a subset, outb, to make sure the mapping works.
+            onnx = ML.Transforms.ApplyOnnxModel(new[] { "outb" }, new[] { "ina", "inb" }, modelFile).Fit(dataView).Transform(dataView);
+
+            outbCol = onnx.Schema["outb"];
+            using (var curs = onnx.GetRowCursor(outbCol))
+            {
+                var getScoresb = curs.GetGetter<VBuffer<float>>(outbCol);
+                var bufferb = default(VBuffer<float>);
+
+                while (curs.MoveNext())
+                {
+                    getScoresb(ref bufferb);
+                    Assert.Equal(5, bufferb.Length);
+                    Assert.Equal(30, bufferb.GetValues().ToArray().Sum());
+                }
+            }
+        }
+
         [OnnxFact]
         public void TestUnknownDimensions()
         {