New ONNX converter interface and some tests

src/Microsoft.ML.Data/Model/ModelOperationsCatalog.cs

@@ -13,6 +13,10 @@ namespace Microsoft.ML
     /// </summary>
     public sealed class ModelOperationsCatalog
     {
+        /// <summary>
+        /// This is a best friend because an extension method defined in another assembly needs this field.
+        /// </summary>
+        [BestFriend]
         internal IHostEnvironment Environment { get; }
 
         public ExplainabilityTransforms Explainability { get; }
@@ -33,7 +37,6 @@ protected SubCatalogBase(ModelOperationsCatalog owner)
             {
                 Environment = owner.Environment;
             }
-
         }
 
         /// <summary>

src/Microsoft.ML.Onnx/OnnxExportExtensions.cs

@@ -0,0 +1,35 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System.Collections.Generic;
+using Microsoft.ML.Core.Data;
+using Microsoft.ML.Data;
+using Microsoft.ML.Model.Onnx;
+using Microsoft.ML.UniversalModelFormat.Onnx;
+
+namespace Microsoft.ML
+{
+    public static class OnnxExportExtensions
+    {
+        /// <summary>
+        /// Convert the specified <see cref="ITransformer"/> to ONNX format. Note that ONNX uses Google's Protobuf so the returned value is a Protobuf object.
+        /// </summary>
+        /// <param name="catalog">The class that <see cref="ConvertToOnnx(ModelOperationsCatalog, ITransformer, IDataView)"/> attached to.</param>
+        /// <param name="transform">The <see cref="ITransformer"/> that will be converted into ONNX format.</param>
+        /// <param name="inputData">The input of the specified transform.</param>
+        /// <returns>An ONNX model equivalent to the converted ML.NET model.</returns>
+        public static ModelProto ConvertToOnnx(this ModelOperationsCatalog catalog, ITransformer transform, IDataView inputData)
+        {
+            var env = catalog.Environment;
+            var ctx = new OnnxContextImpl(env, "model", "ML.NET", "0", 0, "machinelearning.dotnet", OnnxVersion.Stable);
+            var outputData = transform.Transform(inputData);
+            LinkedList<ITransformCanSaveOnnx> transforms = null;
+            using (var ch = env.Start("ONNX conversion"))
+            {
+                SaveOnnxCommand.GetPipe(ctx, ch, outputData, out IDataView root, out IDataView sink, out transforms);
+                return SaveOnnxCommand.ConvertTransformListToOnnxModel(ctx, ch, root, sink, transforms, null, null);
+            }
+        }
+    }
+}

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,10 @@ public override void Run()
             }
         }
 
-        private void GetPipe(OnnxContextImpl ctx, IChannel ch, IDataView end, out IDataView source, out IDataView trueEnd, out LinkedList<ITransformCanSaveOnnx> transforms)
+        internal static void GetPipe(OnnxContextImpl ctx, IChannel ch, IDataView end, out IDataView source, out IDataView trueEnd, out LinkedList<ITransformCanSaveOnnx> transforms)
         {
-            Host.AssertValue(end);
+            ch.AssertValue(end);
+
             source = trueEnd = (end as CompositeDataLoader)?.View ?? end;
             IDataTransform transform = source as IDataTransform;
             transforms = new LinkedList<ITransformCanSaveOnnx>();
@@ -134,7 +136,53 @@ private void GetPipe(OnnxContextImpl ctx, IChannel ch, IDataView end, out IDataV
                 transform = (source = transform.Source) as IDataTransform;
             }
 
-            Host.AssertValue(source);
+            ch.AssertValue(source);
+        }
+
+        internal static ModelProto ConvertTransformListToOnnxModel(OnnxContextImpl ctx, IChannel ch, 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)
+            {
+                ch.Assert(trans.CanSaveOnnx(ctx));
+                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 +258,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, ch, source, end, transforms, _inputsToDrop, _outputsToDrop);
 
-            var model = ctx.MakeModel();
             using (var file = Host.CreateOutputFile(_outputModelPath))
             using (var stream = file.CreateWriteStream())
                 model.WriteTo(stream);