Reuse existing code to do conversion and polish a test

Author: wschin

src/Microsoft.ML.Onnx/SaveOnnxCommand.cs

@@ -12,6 +12,7 @@
 using Microsoft.ML.EntryPoints;
 using Microsoft.ML.Internal.Utilities;
 using Microsoft.ML.Model.Onnx;
+using Microsoft.ML.UniversalModelFormat.Onnx;
 using Newtonsoft.Json;
 
 [assembly: LoadableClass(SaveOnnxCommand.Summary, typeof(SaveOnnxCommand), typeof(SaveOnnxCommand.Arguments), typeof(SignatureCommand),
@@ -113,9 +114,11 @@ public override void Run()
             }
         }
 
-        private void GetPipe(OnnxContextImpl ctx, IChannel ch, IDataView end, out IDataView source, out IDataView trueEnd, out LinkedList<ITransformCanSaveOnnx> transforms)
+        [BestFriend]
+        internal static void GetPipe(OnnxContextImpl ctx, IChannel ch, IDataView end, out IDataView source, out IDataView trueEnd, out LinkedList<ITransformCanSaveOnnx> transforms)
         {
-            Host.AssertValue(end);
+            Contracts.AssertValue(end);
+
             source = trueEnd = (end as CompositeDataLoader)?.View ?? end;
             IDataTransform transform = source as IDataTransform;
             transforms = new LinkedList<ITransformCanSaveOnnx>();
@@ -134,7 +137,51 @@ private void GetPipe(OnnxContextImpl ctx, IChannel ch, IDataView end, out IDataV
                 transform = (source = transform.Source) as IDataTransform;
             }
 
-            Host.AssertValue(source);
+            Contracts.AssertValue(source);
+        }
+
+        [BestFriend]
+        internal static ModelProto ConvertTransformListToOnnxModel(OnnxContextImpl ctx, IDataView inputData, IDataView outputData,
+            LinkedList<ITransformCanSaveOnnx> transforms, HashSet<string> inputColumnNamesToDrop=null, HashSet<string> outputColumnNamesToDrop=null)
+        {
+            inputColumnNamesToDrop = inputColumnNamesToDrop ?? new HashSet<string>();
+            outputColumnNamesToDrop = outputColumnNamesToDrop ?? new HashSet<string>();
+            HashSet<string> inputColumns = new HashSet<string>();
+            // Create graph inputs.
+            for (int i = 0; i < inputData.Schema.Count; i++)
+            {
+                string colName = inputData.Schema[i].Name;
+                if(inputColumnNamesToDrop.Contains(colName))
+                    continue;
+
+                ctx.AddInputVariable(inputData.Schema[i].Type, colName);
+                inputColumns.Add(colName);
+            }
+
+            // Create graph nodes, outputs and intermediate values.
+            foreach (var trans in transforms)
+                trans.SaveAsOnnx(ctx);
+
+            // Add graph outputs.
+            for (int i = 0; i < outputData.Schema.Count; ++i)
+            {
+                if (outputData.Schema[i].IsHidden)
+                    continue;
+
+                var idataviewColumnName = outputData.Schema[i].Name;
+
+                // Since the last IDataView also contains columns of the initial IDataView, last IDataView's columns found in
+                // _inputToDrop should be removed too.
+                if (inputColumnNamesToDrop.Contains(idataviewColumnName) || outputColumnNamesToDrop.Contains(idataviewColumnName))
+                    continue;
+
+                var variableName = ctx.TryGetVariableName(idataviewColumnName);
+                var trueVariableName = ctx.AddIntermediateVariable(null, idataviewColumnName, true);
+                ctx.CreateNode("Identity", variableName, trueVariableName, ctx.GetNodeName("Identity"), "");
+                ctx.AddOutputVariable(outputData.Schema[i].Type, trueVariableName);
+            }
+
+            return ctx.MakeModel();
         }
 
         private void Run(IChannel ch)
@@ -210,45 +257,8 @@ private void Run(IChannel ch)
                     nameof(Arguments.LoadPredictor), "We were explicitly told to load the predictor but one was not present.");
             }
 
-            HashSet<string> inputColumns = new HashSet<string>();
-            //Create graph inputs.
-            for (int i = 0; i < source.Schema.Count; i++)
-            {
-                string colName = source.Schema[i].Name;
-                if(_inputsToDrop.Contains(colName))
-                    continue;
-
-                ctx.AddInputVariable(source.Schema[i].Type, colName);
-                inputColumns.Add(colName);
-            }
-
-            //Create graph nodes, outputs and intermediate values.
-            foreach (var trans in transforms)
-            {
-                Host.Assert(trans.CanSaveOnnx(ctx));
-                trans.SaveAsOnnx(ctx);
-            }
-
-            //Add graph outputs.
-            for (int i = 0; i < end.Schema.Count; ++i)
-            {
-                if (end.Schema[i].IsHidden)
-                    continue;
-
-                var idataviewColumnName = end.Schema[i].Name;
-
-                // Since the last IDataView also contains columns of the initial IDataView, last IDataView's columns found in
-                // _inputToDrop should be removed too.
-                if (_inputsToDrop.Contains(idataviewColumnName) || _outputsToDrop.Contains(idataviewColumnName))
-                    continue;
-
-                var variableName = ctx.TryGetVariableName(idataviewColumnName);
-                var trueVariableName = ctx.AddIntermediateVariable(null, idataviewColumnName, true);
-                ctx.CreateNode("Identity", variableName, trueVariableName, ctx.GetNodeName("Identity"), "");
-                ctx.AddOutputVariable(end.Schema[i].Type, trueVariableName);
-            }
+            var model = ConvertTransformListToOnnxModel(ctx, source, end, transforms, _inputsToDrop, _outputsToDrop);
 
-            var model = ctx.MakeModel();
             using (var file = Host.CreateOutputFile(_outputModelPath))
             using (var stream = file.CreateWriteStream())
                 model.WriteTo(stream);

src/Microsoft.ML.Onnx/TransformerChainOnnxConverter.cs

@@ -1,42 +1,24 @@
-using Microsoft.ML.Core.Data;
+using System.Collections.Generic;
+using Microsoft.ML.Core.Data;
 using Microsoft.ML.Data;
 using Microsoft.ML.UniversalModelFormat.Onnx;
 
 namespace Microsoft.ML.Model.Onnx
 {
     public class TransformerChainOnnxConverter
     {
-        public static ModelProto Convert<T>(TransformerChain<T> chain, Schema inputSchema) where T : class, ITransformer
+        public static ModelProto Convert<T>(TransformerChain<T> chain, IDataView inputData, HashSet<string> inputColumnNamesToDrop=null, HashSet<string> outputColumnNamesToDrop=null) where T : class, ITransformer
         {
             var env = new MLContext();
-            var onnxContext = new OnnxContextImpl(env, "model", "ML.NET", "0", 0, "com.test", Model.Onnx.OnnxVersion.Stable);
+            var ctx = new OnnxContextImpl(env, "model", "ML.NET", "0", 0, "com.test", OnnxVersion.Stable);
+            var outputData = chain.Transform(inputData);
+            IDataView source = null;
+            IDataView trueEnd = null;
+            LinkedList<ITransformCanSaveOnnx> transforms = null;
+            using (var ch = (env as IChannelProvider).Start("ONNX conversion"))
+                SaveOnnxCommand.GetPipe(ctx, ch, outputData, out source, out trueEnd, out transforms);
 
-            for (int i = 0; i < inputSchema.Count; i++)
-            {
-                string colName = inputSchema[i].Name;
-                onnxContext.AddInputVariable(inputSchema[i].Type, colName);
-            }
-
-            foreach (var t in chain)
-            {
-                var mapper = t.GetRowToRowMapper(inputSchema);
-                inputSchema = t.GetOutputSchema(inputSchema);
-                (mapper as ISaveAsOnnx).SaveAsOnnx(onnxContext);
-            }
-
-            for (int i = 0; i < inputSchema.Count; ++i)
-            {
-                if (inputSchema[i].IsHidden)
-                    continue;
-
-                var idataviewColumnName = inputSchema[i].Name;
-
-                var variableName = onnxContext.TryGetVariableName(idataviewColumnName);
-                var trueVariableName = onnxContext.AddIntermediateVariable(null, idataviewColumnName, true);
-                onnxContext.CreateNode("Identity", variableName, trueVariableName, onnxContext.GetNodeName("Identity"), "");
-                onnxContext.AddOutputVariable(inputSchema[i].Type, trueVariableName);
-            }
-            return onnxContext.MakeModel();
+            return SaveOnnxCommand.ConvertTransformListToOnnxModel(ctx, source, trueEnd, transforms, inputColumnNamesToDrop, outputColumnNamesToDrop);
         }
     }
 }

test/Microsoft.ML.Tests/OnnxConversionTest.cs

@@ -1,5 +1,4 @@
-using System.Collections.Generic;
-using System.Linq;
+using System.Linq;
 using Google.Protobuf;
 using Microsoft.ML.Data;
 using Microsoft.ML.Model.Onnx;
@@ -25,41 +24,45 @@ public OnnxConversionTest(ITestOutputHelper output) : base(output)
         {
         }
 
+        /// <summary>
+        /// In this test, we convert a trained <see cref="TransformerChain"/> into ONNX <see cref="UniversalModelFormat.Onnx.ModelProto"/> file and then
+        /// call <see cref="OnnxScoringEstimator"/> to evaluate that file. The outputs of <see cref="OnnxScoringEstimator"/> are checked against the original
+        /// ML.NET model's outputs. 
+        /// </summary>
         [Fact]
-        public void SimplePipelineOnnxConversionTest()
+        public void SimpleEndToEndOnnxConversionTest()
         {
+            // Step 1: Create and train a ML.NET pipeline.
             var trainDataPath = GetDataPath(TestDatasets.generatedRegressionDataset.trainFilename);
             var mlContext = new MLContext();
-
-            var trainData = mlContext.Data.ReadFromTextFile<AdultData>(trainDataPath,
+            var data = mlContext.Data.ReadFromTextFile<AdultData>(trainDataPath,
                 hasHeader: true,
                 separatorChar: ';'
             );
-
-            var cachedTrainData = mlContext.Data.Cache(trainData);
-
+            var cachedTrainData = mlContext.Data.Cache(data);
             var dynamicPipeline =
                 mlContext.Transforms.Normalize("FeatureVector")
                 .AppendCacheCheckpoint(mlContext)
                 .Append(mlContext.Regression.Trainers.StochasticDualCoordinateAscent(labelColumn: "Target", featureColumn: "FeatureVector"));
+            var model = dynamicPipeline.Fit(data);
+            var transformedData = model.Transform(data);
 
-            var model = dynamicPipeline.Fit(trainData);
-            var transformedData = model.Transform(trainData);
-
-            var onnxModel = TransformerChainOnnxConverter.Convert(model, trainData.Schema);
-
+            // Step 2: Convert ML.NET model to ONNX format and save it as a file.
+            var onnxModel = TransformerChainOnnxConverter.Convert(model, data);
             var onnxFileName = "model.onnx";
             var onnxFilePath = GetOutputPath(onnxFileName);
             using (var file = (mlContext as IHostEnvironment).CreateOutputFile(onnxFilePath))
             using (var stream = file.CreateWriteStream())
                 onnxModel.WriteTo(stream);
 
+            // Step 3: Evaluate the saved ONNX model using the data used to train the ML.NET pipeline.
             string[] inputNames = onnxModel.Graph.Input.Select(valueInfoProto => valueInfoProto.Name).ToArray();
             string[] outputNames = onnxModel.Graph.Output.Select(valueInfoProto => valueInfoProto.Name).ToArray();
             var onnxEstimator = new OnnxScoringEstimator(mlContext, onnxFilePath, inputNames, outputNames);
-            var onnxTransformer = onnxEstimator.Fit(trainData);
-            var onnxResult = onnxTransformer.Transform(trainData);
+            var onnxTransformer = onnxEstimator.Fit(data);
+            var onnxResult = onnxTransformer.Transform(data);
 
+            // Step 4: Compare ONNX and ML.NET results.
             using (var expectedCursor = transformedData.GetRowCursor(columnIndex => columnIndex == transformedData.Schema["Score"].Index))
             using (var actualCursor = onnxResult.GetRowCursor(columnIndex => columnIndex == onnxResult.Schema["Score0"].Index))
             {