From 187a42bb6481a9f6fa1599dfe54c63ade0b95b21 Mon Sep 17 00:00:00 2001 From: "Harish S. Kulkarni" Date: Sun, 23 Feb 2020 22:25:57 -0800 Subject: [PATCH] Slightly simplified version of adding KeyToValue for onnx export --- .../SaveOnnxCommand.cs | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/src/Microsoft.ML.OnnxConverter/SaveOnnxCommand.cs b/src/Microsoft.ML.OnnxConverter/SaveOnnxCommand.cs index d29464d5fc..295da4bddc 100644 --- a/src/Microsoft.ML.OnnxConverter/SaveOnnxCommand.cs +++ b/src/Microsoft.ML.OnnxConverter/SaveOnnxCommand.cs @@ -15,6 +15,7 @@ using Microsoft.ML.Internal.Utilities; using Microsoft.ML.Model.OnnxConverter; using Microsoft.ML.Runtime; +using Microsoft.ML.Transforms; using Newtonsoft.Json; using static Microsoft.ML.Model.OnnxConverter.OnnxCSharpToProtoWrapper; @@ -322,6 +323,26 @@ private void Run(IChannel ch) nameof(Arguments.LoadPredictor), "We were explicitly told to load the predictor but one was not present."); } + // If the output schema has KeyDataViewType columns, and it has annotations + // add KeyToValueMappingTransformer to translate the keys back to values + // This is done during Onnx export but not in regular pipelines because onnx does not + // support annotations on tensors + foreach (var column in end.Schema) + { + if ((column.IsHidden) || (!(column.Type.GetItemType() is KeyDataViewType keyType))) + continue; + + var metaColumn = column.Annotations.Schema.GetColumnOrNull(AnnotationUtils.Kinds.KeyValues); + if (metaColumn != null + && metaColumn.Value.Type is VectorDataViewType vectorType + && keyType.Count == (ulong)vectorType.Size) + { + var outputData = new KeyToValueMappingTransformer(Host, column.Name).Transform(end); + end = outputData; + transforms.AddLast(end as ITransformCanSaveOnnx); + } + } + var model = ConvertTransformListToOnnxModel(ctx, ch, source, end, transforms, _inputsToDrop, _outputsToDrop); using (var file = Host.CreateOutputFile(_outputModelPath))