ONNXTransform Upgrade to Enable Non-tensor Types

build/Dependencies.props

@@ -44,9 +44,9 @@
   <PropertyGroup>
     <BenchmarkDotNetVersion>0.11.3</BenchmarkDotNetVersion>
     <MicrosoftCodeAnalysisTestingVersion>1.0.0-beta1-63812-02</MicrosoftCodeAnalysisTestingVersion>
-    <MicrosoftMLTestModelsPackageVersion>0.0.4-test</MicrosoftMLTestModelsPackageVersion>
+    <MicrosoftMLTestModelsPackageVersion>0.0.5-test</MicrosoftMLTestModelsPackageVersion>
     <MicrosoftMLTensorFlowTestModelsVersion>0.0.11-test</MicrosoftMLTensorFlowTestModelsVersion>
-    <MicrosoftMLOnnxTestModelsVersion>0.0.4-test</MicrosoftMLOnnxTestModelsVersion>
+    <MicrosoftMLOnnxTestModelsVersion>0.0.5-test</MicrosoftMLOnnxTestModelsVersion>
   </PropertyGroup>
 
 </Project>

src/Microsoft.ML.OnnxConverter/AssemblyInfo.cs

@@ -7,3 +7,4 @@
 
 [assembly: InternalsVisibleTo("Microsoft.ML.Core.Tests" + PublicKey.TestValue)]
 [assembly: InternalsVisibleTo("Microsoft.ML.Tests" + PublicKey.TestValue)]
+[assembly: InternalsVisibleTo("Microsoft.ML.OnnxTransformer" + PublicKey.Value)]

src/Microsoft.ML.OnnxTransformer/Microsoft.ML.OnnxTransformer.csproj

@@ -10,6 +10,7 @@
     <ProjectReference Include="..\Microsoft.ML.Core\Microsoft.ML.Core.csproj" />
     <ProjectReference Include="..\Microsoft.ML.Data\Microsoft.ML.Data.csproj" />
     <PackageReference Include="Microsoft.ML.OnnxRuntime" Version="$(MicrosoftMLOnnxRuntimePackageVersion)" />
+    <PackageReference Include="Google.Protobuf" Version="$(GoogleProtobufPackageVersion)" />
   </ItemGroup>
 
 </Project>

src/Microsoft.ML.OnnxTransformer/OnnxUtils.cs

@@ -8,6 +8,7 @@
 using System.Linq;
 using System.Numerics.Tensors;
 using Microsoft.ML.Data;
+using Microsoft.ML.Model.OnnxConverter;
 using Microsoft.ML.OnnxRuntime;
 using Microsoft.ML.Runtime;
 using OnnxShape = System.Collections.Generic.List<int>;
@@ -20,67 +21,137 @@ namespace Microsoft.ML.Transforms.Onnx
     /// It provides API to open a session, score tensors (NamedOnnxValues) and return
     /// the results.
     /// </summary>
-    internal sealed class OnnxModel
+    internal sealed class OnnxModel : IDisposable
     {
-
         /// <summary>
         /// OnnxModelInfo contains the data that we should get from
         /// OnnxRuntime API once that functionality is added.
         /// </summary>
         public sealed class OnnxModelInfo
         {
-            public readonly OnnxNodeInfo[] InputsInfo;
-            public readonly OnnxNodeInfo[] OutputsInfo;
+            /// <summary>
+            /// InputNames[i] is the name of the i-th element in <see cref="InputsInfo"/>.
+            /// </summary>
+            public List<string> InputNames { get; }
+            /// <summary>
+            /// OutputNames[i] is the name of the i-th element in <see cref="OutputsInfo"/>.
+            /// </summary>
+            public List<string> OutputNames { get; }
+            /// <summary>
+            /// Inputs of the containing <see cref="OnnxModel"/>.
+            /// </summary>
+            public OnnxVariableInfo[] InputsInfo { get; }
+            /// <summary>
+            /// Outputs of the containing <see cref="OnnxModel"/>.
+            /// </summary>
+            public OnnxVariableInfo[] OutputsInfo { get; }
 
-            public OnnxModelInfo(IEnumerable<OnnxNodeInfo> inputsInfo, IEnumerable<OnnxNodeInfo> outputsInfo)
+            public OnnxModelInfo(IEnumerable<OnnxVariableInfo> inputsInfo, IEnumerable<OnnxVariableInfo> outputsInfo)
             {
+                InputNames = inputsInfo.Select(val => val.Name).ToList();
                 InputsInfo = inputsInfo.ToArray();
+                OutputNames = outputsInfo.Select(val => val.Name).ToList();
                 OutputsInfo = outputsInfo.ToArray();
             }
+
+            /// <summary>
+            /// Return the ONNX value for a <see cref="IDataView"/> input column called <paramref name="name"/>.
+            /// </summary>
+            public OnnxVariableInfo GetInput(string name)
+            {
+                var index = InputNames.IndexOf(name);
+                if (index < 0)
+                    throw Contracts.ExceptParamValue(name, nameof(name), $"Input tensor, {name}, does not exist in the ONNX model. " +
+                        $"Available input names are [{string.Join(",", InputNames)}].");
+                return InputsInfo[index];
+            }
+
+            /// <summary>
+            /// Return the ONNX value for a <see cref="IDataView"/> output column called <paramref name="name"/>.
+            /// </summary>
+            public OnnxVariableInfo GetOutput(string name)
+            {
+                var index = OutputNames.IndexOf(name);
+                if (index < 0)
+                    throw Contracts.ExceptParamValue(name, nameof(name), $"Onput tensor, {name}, does not exist in the ONNX model. " +
+                        $"Available output names are [{string.Join(",", OutputNames)}].");
+                return OutputsInfo[index];
+            }
         }
 
         /// <summary>
         /// OnnxNodeInfo contains all the information for a given node (e.g. inputs/outputs)
         /// of an Onnx model.
         /// </summary>
-        public class OnnxNodeInfo
+        public class OnnxVariableInfo
         {
             /// <summary>
-            /// The Name of the node
+            /// The Name of the variable. Note that ONNX variable are named.
             /// </summary>
-            public readonly string Name;
+            public string Name { get; }
             /// <summary>
-            /// The shape of the node
+            /// The shape of the variable if the variable is a tensor. For other
+            /// types such sequence and dictionary, <see cref="Shape"/> would be
+            /// <see langword="null"/>.
             /// </summary>
-            public readonly OnnxShape Shape;
+            public OnnxShape Shape { get; }
             /// <summary>
-            /// The type of the node
+            /// The type of the variable produced by ONNXRuntime.
             /// </summary>
-            public readonly System.Type Type;
+            public Type TypeInOnnxRuntime { get; }
+            /// <summary>
+            /// The <see cref="Data.DataViewType"/> that this ONNX variable corresponds
+            /// to in <see cref="IDataView"/>'s type system.
+            /// </summary>
+            public DataViewType DataViewType { get; }
+            /// <summary>
+            /// A method to case <see cref="NamedOnnxValue"/> produced by
+            /// ONNXRuntime to the type specified in <see cref="DataViewType"/>.
+            /// </summary>
+            public Func<NamedOnnxValue, object> Caster { get; }
 
-            public OnnxNodeInfo(string name, OnnxShape shape, System.Type type)
+            public OnnxVariableInfo(string name, OnnxShape shape, Type typeInOnnxRuntime, DataViewType mlnetType, Func<NamedOnnxValue, object> caster)
             {
                 Name = name;
                 Shape = shape;
-                Type = type;
+                TypeInOnnxRuntime = typeInOnnxRuntime;
+                DataViewType = mlnetType;
+                Caster = caster;
             }
         }
 
-        public readonly OnnxModelInfo ModelInfo;
+        /// <summary>
+        /// The ONNXRuntime facility to execute the loaded ONNX model.
+        /// </summary>
         private readonly InferenceSession _session;
-        private readonly string _modelFile;
-        public readonly List<string> InputNames;
-        public readonly List<string> OutputNames;
+        /// <summary>
+        /// Indicates if <see cref="ModelFile"/> is a temporal file created by <see cref="CreateFromBytes(byte[], int?, bool)"/>
+        /// or <see cref="CreateFromBytes(byte[])"/>. If <see langword="true"/>, <see cref="Dispose(bool)"/> should delete <see cref="ModelFile"/>.
+        /// </summary>
+        private bool _ownModelFile;
+        /// <summary>
+        /// The location where the used ONNX model loaded from.
+        /// </summary>
+        internal string ModelFile { get; }
+        /// <summary>
+        /// The ONNX model file that <see cref="OnnxModel"/> built upon.
+        /// </summary>
+        internal OnnxModelInfo ModelInfo { get; }
 
         /// <summary>
         /// Constructs OnnxModel object from file.
         /// </summary>
         /// <param name="modelFile">Model file path.</param>
         /// <param name="gpuDeviceId">GPU device ID to execute on. Null for CPU.</param>
         /// <param name="fallbackToCpu">If true, resumes CPU execution quitely upon GPU error.</param>
-        public OnnxModel(string modelFile, int? gpuDeviceId = null, bool fallbackToCpu = false)
+        /// <param name="ownModelFile">If true, the <paramref name="modelFile"/> will be deleted when <see cref="OnnxModel"/> is
+        /// no longer needed.</param>
+        public OnnxModel(string modelFile, int? gpuDeviceId = null, bool fallbackToCpu = false, bool ownModelFile=false)
         {
-            _modelFile = modelFile;
+            ModelFile = modelFile;
+            // If we don't own the model file, _disposed should be false to prevent deleting user's file.
+            _ownModelFile = ownModelFile;
+            _disposed = false;
 
             if (gpuDeviceId != null)
             {
@@ -103,13 +174,55 @@ public OnnxModel(string modelFile, int? gpuDeviceId = null, bool fallbackToCpu =
                 _session = new InferenceSession(modelFile);
             }
 
-            ModelInfo = new OnnxModelInfo(GetInputsInfo(), GetOutputsInfo());
-            InputNames = ModelInfo.InputsInfo.Select(i => i.Name).ToList();
-            OutputNames = ModelInfo.OutputsInfo.Select(i => i.Name).ToList();
+            // Load ONNX model file and parse its input and output schema. The reason of doing so is that ONNXRuntime
+            // doesn't expose full type information via its C# APIs.
+            ModelFile = modelFile;
+            var model = new OnnxCSharpToProtoWrapper.ModelProto();
+            using (var modelStream = File.OpenRead(modelFile))
+                model = OnnxCSharpToProtoWrapper.ModelProto.Parser.ParseFrom(modelStream);
+
+            // Parse actual input and output types stored in the loaded ONNX model to get their DataViewType's.
+            var inputTypePool = new Dictionary<string, DataViewType>();
+            foreach (var valueInfo in model.Graph.Input)
+                inputTypePool[valueInfo.Name] = OnnxTypeHelper.GetDataViewType(valueInfo.Type);
+            var outputTypePool = new Dictionary<string, DataViewType>();
+
+            // Build casters which maps NamedOnnxValue to .NET objects.
+            var casterPool = new Dictionary<string, Func<NamedOnnxValue, object>>();
+            foreach (var valueInfo in model.Graph.Output)
+            {
+                outputTypePool[valueInfo.Name] = OnnxTypeHelper.GetDataViewType(valueInfo.Type);
+                casterPool[valueInfo.Name] = OnnxTypeHelper.GetDataViewValueCasterAndResultedType(valueInfo.Type, out Type actualType);
+            }
+
+            var onnxRuntimeInputInfos = new List<OnnxVariableInfo>();
+            foreach (var pair in _session.InputMetadata)
+            {
+                var name = pair.Key;
+                var meta = pair.Value;
+                var dataViewType = inputTypePool[name];
+                var info = new OnnxVariableInfo(name, meta.Dimensions.ToList(), meta.ElementType, dataViewType, null);
+                onnxRuntimeInputInfos.Add(info);
+            }
+
+            var onnxRuntimeOutputInfos = new List<OnnxVariableInfo>();
+            foreach (var pair in _session.OutputMetadata)
+            {
+                var name = pair.Key;
+                var meta = pair.Value;
+                var dataViewType = outputTypePool[name];
+                var caster = casterPool[name];
+                var info = new OnnxVariableInfo(name, meta.Dimensions.ToList(), meta.ElementType, dataViewType, caster);
+                onnxRuntimeOutputInfos.Add(info);
+            }
+
+            ModelInfo = new OnnxModelInfo(onnxRuntimeInputInfos, onnxRuntimeOutputInfos);
         }
 
         /// <summary>
-        /// Create an OnnxModel from a byte[]
+        /// Create an OnnxModel from a byte[]. Usually, a ONNX model is consumed by <see cref="OnnxModel"/> as a file.
+        /// With <see cref="CreateFromBytes(byte[])"/> and <see cref="CreateFromBytes(byte[], int?, bool)"/>, it's possible
+        /// to use in-memory model (type: byte[]) to create <see cref="OnnxModel"/>.
         /// </summary>
         /// <param name="modelBytes">Bytes of the serialized model</param>
         /// <returns>OnnxModel</returns>
@@ -120,6 +233,9 @@ public static OnnxModel CreateFromBytes(byte[] modelBytes)
 
         /// <summary>
         /// Create an OnnxModel from a byte[]. Set execution to GPU if required.
+        /// Usually, a ONNX model is consumed by <see cref="OnnxModel"/> as a file.
+        /// With <see cref="CreateFromBytes(byte[])"/> and <see cref="CreateFromBytes(byte[], int?, bool)"/>,
+        /// it's possible to use in-memory model (type: byte[]) to create <see cref="OnnxModel"/>.
         /// </summary>
         /// <param name="modelBytes">Bytes of the serialized model.</param>
         /// <param name="gpuDeviceId">GPU device ID to execute on. Null for CPU.</param>
@@ -132,12 +248,7 @@ public static OnnxModel CreateFromBytes(byte[] modelBytes, int? gpuDeviceId = nu
 
             var tempModelFile = Path.Combine(tempModelDir, "model.onnx");
             File.WriteAllBytes(tempModelFile, modelBytes);
-            return new OnnxModel(tempModelFile, gpuDeviceId, fallbackToCpu);
-
-            // TODO:
-            // tempModelFile is needed in case the model needs to be saved
-            // Either have to save the modelbytes and delete the temp dir/file,
-            // or keep the dir/file and write proper cleanup when application closes
+            return new OnnxModel(tempModelFile, gpuDeviceId, fallbackToCpu, ownModelFile: true);
         }
 
         /// <summary>
@@ -151,37 +262,49 @@ public IReadOnlyCollection<NamedOnnxValue> Run(List<NamedOnnxValue> inputNamedOn
         }
 
         /// <summary>
-        /// Convert the model to a byte array.
+        /// Flag used to indicate if the unmanaged resources (aka the model file <see cref="ModelFile"/>
+        /// and <see cref="_session"/>) have been deleted.
         /// </summary>
-        /// <returns>byte[]</returns>
-        public byte[] ToByteArray()
+        private bool _disposed;
+
+        public void Dispose()
         {
-            return File.ReadAllBytes(_modelFile);
+            Dispose(true);
+            GC.SuppressFinalize(this);
         }
 
         /// <summary>
-        /// Returns input metadata of the ONNX model.
+        /// There are two unmanaged resources we can dispose, <see cref="_session"/> and <see cref="ModelFile"/>
+        /// if <see cref="_ownModelFile"/> is <see langword="true"/>.
         /// </summary>
-        /// <returns>OnnxNodeInfo[]</returns>
-        private IEnumerable<OnnxNodeInfo> GetInputsInfo()
+        /// <param name="disposing"></param>
+        private void Dispose(bool disposing)
         {
-            return _session.InputMetadata.Select(kv => new OnnxNodeInfo(kv.Key, kv.Value.Dimensions.ToList(), kv.Value.ElementType));
+            if (!_disposed)
+            {
+                // There are two things to be disposed.
+                if (disposing)
+                {
+                    // First, we release the resource token by ONNXRuntime.
+                    _session.Dispose();
+                    // Second, we delete the model file if that file is not created by the user.
+                    if (_ownModelFile && File.Exists(ModelFile))
+                        File.Delete(ModelFile);
+                }
+                _disposed = true;
+            }
         }
 
-        /// <summary>
-        /// Returns output metadata of the ONNX model.
-        /// </summary>
-        /// <returns></returns>
-        private IEnumerable<OnnxNodeInfo> GetOutputsInfo()
+        ~OnnxModel()
         {
-            return _session.OutputMetadata.Select(kv => new OnnxNodeInfo(kv.Key, kv.Value.Dimensions.ToList(), kv.Value.ElementType));
+            Dispose(false);
         }
     }
 
     internal sealed class OnnxUtils
     {
-        private static HashSet<System.Type> _onnxTypeMap =
-            new HashSet<System.Type>
+        private static HashSet<Type> _onnxTypeMap =
+            new HashSet<Type>
                 {
                      typeof(Double),
                      typeof(Single),
@@ -192,8 +315,8 @@ internal sealed class OnnxUtils
                      typeof(UInt32),
                      typeof(UInt64)
                 };
-        private static Dictionary<System.Type, InternalDataKind> _typeToKindMap=
-            new Dictionary<System.Type, InternalDataKind>
+        private static Dictionary<Type, InternalDataKind> _typeToKindMap=
+            new Dictionary<Type, InternalDataKind>
                 {
                     { typeof(Single) , InternalDataKind.R4},
                     { typeof(Double) , InternalDataKind.R8},
@@ -243,7 +366,7 @@ public static NamedOnnxValue CreateNamedOnnxValue<T>(string name, ReadOnlySpan<T
         /// </summary>
         /// <param name="type"></param>
         /// <returns></returns>
-        public static PrimitiveDataViewType OnnxToMlNetType(System.Type type)
+        public static PrimitiveDataViewType OnnxToMlNetType(Type type)
         {
             if (!_typeToKindMap.ContainsKey(type))
                throw Contracts.ExceptNotSupp("Onnx type not supported", type);