diff --git a/src/Microsoft.ML.OnnxTransformer/OnnxTransform.cs b/src/Microsoft.ML.OnnxTransformer/OnnxTransform.cs index 65efbaf419..60d82c6a76 100644 --- a/src/Microsoft.ML.OnnxTransformer/OnnxTransform.cs +++ b/src/Microsoft.ML.OnnxTransformer/OnnxTransform.cs @@ -389,7 +389,14 @@ public Mapper(OnnxTransformer parent, DataViewSchema inputSchema) : throw Host.Except($"Variable length input columns not supported"); if (type.GetItemType() != inputNodeInfo.DataViewType.GetItemType()) - throw Host.ExceptSchemaMismatch(nameof(inputSchema), "input", _parent.Inputs[i], inputNodeInfo.DataViewType.GetItemType().ToString(), type.ToString()); + { + // If the ONNX model input node expects a type that mismatches with the type of the input IDataView column that is provided + // then throw an exception. + // This is done except in the case where the ONNX model input node expects a UInt32 but the input column is actually KeyDataViewType + // This is done to support a corner case originated in NimbusML. For more info, see: https://github.com/microsoft/NimbusML/issues/426 + if (!(type.GetItemType() is KeyDataViewType && inputNodeInfo.DataViewType.GetItemType().RawType == typeof(UInt32))) + throw Host.ExceptSchemaMismatch(nameof(inputSchema), "input", _parent.Inputs[i], inputNodeInfo.DataViewType.GetItemType().ToString(), type.ToString()); + } // 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. diff --git a/test/Microsoft.ML.Tests/OnnxConversionTest.cs b/test/Microsoft.ML.Tests/OnnxConversionTest.cs index 75f5b52466..8e77b8bf15 100644 --- a/test/Microsoft.ML.Tests/OnnxConversionTest.cs +++ b/test/Microsoft.ML.Tests/OnnxConversionTest.cs @@ -1450,6 +1450,81 @@ public void CopyColumnsOnnxTest() Done(); } + [Fact] + public void UseKeyDataViewTypeAsUInt32InOnnxInput() + { + // In this test an onnx model which expect a uin32 input column is applied to a KeyDataViewType input column + // This, is done as needed by NimbusML. For more context see: https://github.com/microsoft/NimbusML/issues/426 + + // Step 1: Load the Iris Dataset and apply a Value To Key Mapping to it. + // Save the resulting dataview in .idv format eliminating all hidden columns + var mlContext = new MLContext(); + var loader = mlContext.Data.CreateTextLoader( + columns: new[] + { + new TextLoader.Column("Label", DataKind.String, 0), + new TextLoader.Column("SepalLength", DataKind.Single, 1), + new TextLoader.Column("SepalWidth", DataKind.Single, 2), + new TextLoader.Column("PetalLength", DataKind.Single, 3), + new TextLoader.Column("PetalWidth", DataKind.Single, 4) + }, + hasHeader: false + ); + + string dataPath = GetDataPath("iris.txt"); + var originalData = loader.Load(dataPath); + var pipeline1 = mlContext.Transforms.Conversion.MapValueToKey("Label"); + var mappedData = pipeline1.Fit(originalData).Transform(originalData); + + string mappedDataPath = GetOutputPath("kdvt-as-uint32-mapped-data.idv"); + using (FileStream stream = new FileStream(mappedDataPath, FileMode.Create)) + mlContext.Data.SaveAsBinary(mappedData, stream, keepHidden: false); + + // Step 2: Load back the saved .idv + // This IDataView will have a Label column of type KeyDataViewType + // It's necessary to do this, because if I were to use mappedData directly inside the next + // steps, then when saving the ONNX model, it would actually also save the ValueToKeyTransformer part + // and that wouldn't reproduce the scenario. + IDataView reloadedData = mlContext.Data.LoadFromBinary(mappedDataPath); + + // Step 3: Create ONNX model which simply applies Identity to Label column + var pipeline2 = mlContext.Transforms.CopyColumns("Label", "Label"); + var model = pipeline2.Fit(reloadedData); + + var onnxModelPath = GetOutputPath("onnxmodel1-kdvt-as-uint32.onnx"); + using (FileStream stream = new FileStream(onnxModelPath, FileMode.Create)) + mlContext.Model.ConvertToOnnx(model, reloadedData, stream); + + // Step 4: Get input and output names of model + var onnxProtoBufModel = mlContext.Model.ConvertToOnnxProtobuf(model, reloadedData); + string[] inputNames = onnxProtoBufModel.Graph.Input.Select(valueInfoProto => valueInfoProto.Name).ToArray(); + string[] outputNames = onnxProtoBufModel.Graph.Output.Select(valueInfoProto => valueInfoProto.Name).ToArray(); + + if (IsOnnxRuntimeSupported()) + { + // Step 5: Apply Onnx Model + var onnxEstimator = mlContext.Transforms.ApplyOnnxModel(outputNames, inputNames, onnxModelPath); + var onnxResult = onnxEstimator.Fit(reloadedData).Transform(reloadedData); + + // Step 6: Compare results to an onnx model created using the mappedData IDataView + // Notice that this ONNX model would actually include the steps to do the ValueToKeyTransformer mapping, + // because mappedData actually includes the information to do the mapping, and so ONNX would that automatically. + // And because of this, it can only be applied to originalData dataview, despite mappedData was used to create the model. + // If it's tried to apply this model to mappedData or reloadedData, it will throw an exception, since the ONNX model + // will expect a Label input of type string (which only originalData provides). + string onnxModelPath2 = GetOutputPath("onnxmodel2-kdvt-as-uint32.onnx"); + using (FileStream stream = new FileStream(onnxModelPath2, FileMode.Create)) + mlContext.Model.ConvertToOnnx(model, mappedData, stream); + var onnxEstimator2 = mlContext.Transforms.ApplyOnnxModel(outputNames, inputNames, onnxModelPath2); + var onnxResult2 = onnxEstimator2.Fit(originalData).Transform(originalData); + + foreach (var name in outputNames) + CompareResults(name, name, onnxResult, onnxResult2); + } + + Done(); + } + [Fact] public void FeatureSelectionOnnxTest() {