Onnxtransform - api changes for GPU support (#1922)

* API changes for GPU support

* update to test package on nuget.org

* Addressed PR comments

* Addressed PR comments

* Point to nuget on myget.org

* restore directory.build.props

Author: jignparm

build/Dependencies.props

@@ -15,7 +15,7 @@
   <PropertyGroup>
     <GoogleProtobufPackageVersion>3.5.1</GoogleProtobufPackageVersion>
     <LightGBMPackageVersion>2.2.1.1</LightGBMPackageVersion>
-    <MicrosoftMLOnnxRuntimePackageVersion>0.1.5</MicrosoftMLOnnxRuntimePackageVersion>
+    <MicrosoftMLOnnxRuntimeGpuPackageVersion>0.1.5</MicrosoftMLOnnxRuntimeGpuPackageVersion>
     <MlNetMklDepsPackageVersion>0.0.0.7</MlNetMklDepsPackageVersion>
     <ParquetDotNetPackageVersion>2.1.3</ParquetDotNetPackageVersion>
     <SystemDrawingCommonPackageVersion>4.5.0</SystemDrawingCommonPackageVersion>

pkg/Microsoft.ML.OnnxTransform/Microsoft.ML.OnnxTransform.nupkgproj

@@ -7,7 +7,7 @@
 
   <ItemGroup>
     <ProjectReference Include="../Microsoft.ML/Microsoft.ML.nupkgproj" />
-    <PackageReference Include="Microsoft.ML.OnnxRuntime" Version="$(MicrosoftMLOnnxRuntimePackageVersion)"/>
+    <PackageReference Include="Microsoft.ML.OnnxRuntime.Gpu" Version="$(MicrosoftMLOnnxRuntimeGpuPackageVersion)"/>
   </ItemGroup>
 
 </Project>

src/Microsoft.ML.OnnxTransform/Microsoft.ML.OnnxTransform.csproj

@@ -9,7 +9,7 @@
   <ItemGroup>
     <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="Microsoft.ML.OnnxRuntime.Gpu" Version="$(MicrosoftMLOnnxRuntimeGpuPackageVersion)" />
   </ItemGroup>
 
 </Project>

src/Microsoft.ML.OnnxTransform/OnnxTransform.cs

@@ -43,9 +43,15 @@ namespace Microsoft.ML.Transforms
     /// </summary>
     /// <remarks>
     /// <p>Supports inferencing of models in 1.2 and 1.3 format, using the
-    /// <a href='https://www.nuget.org/packages/Microsoft.ML.OnnxRuntime/'>Microsoft.ML.OnnxRuntime</a> library
+    /// <a href='https://www.nuget.org/packages/Microsoft.ML.OnnxRuntime/'>Microsoft.ML.OnnxRuntime</a> library.
     /// </p>
-    /// <p>The inputs and outputs of the onnx models must of of Tensors. Sequence and Maps are not yet supported.</p>
+    /// <p>Models are scored on CPU by default. If GPU execution is needed (optional), install
+    /// <a href='https://developer.nvidia.com/cuda-downloads'>CUDA 10.0 Toolkit</a>
+    /// and
+    /// <a href='https://developer.nvidia.com/cudnn'>cuDNN</a>
+    /// , and set the parameter 'gpuDeviceId' to a valid non-negative integer. Typical device ID values are 0 or 1.
+    /// </p>
+    /// <p>The inputs and outputs of the ONNX models must be Tensor type. Sequence and Maps are not yet supported.</p>
     /// <p>Visit https://github.com/onnx/models to see a list of readily available models to get started with.</p>
     /// <p>Refer to http://onnx.ai' for more information about ONNX.</p>
     /// </remarks>
@@ -61,6 +67,12 @@ public sealed class Arguments : TransformInputBase
 
             [Argument(ArgumentType.Multiple | ArgumentType.Required, HelpText = "Name of the output column.", SortOrder = 2)]
             public string[] OutputColumns;
+
+            [Argument(ArgumentType.AtMostOnce | ArgumentType.Required, HelpText = "GPU device id to run on (e.g. 0,1,..). Null for CPU. Requires CUDA 10.0.", SortOrder = 3)]
+            public int? GpuDeviceId = null;
+
+            [Argument(ArgumentType.AtMostOnce | ArgumentType.Required, HelpText = "If true, resumes execution on CPU upon GPU error. If false, will raise the GPU execption.", SortOrder = 4)]
+            public bool FallbackToCpu = false;
         }
 
         private readonly Arguments _args;
@@ -88,15 +100,27 @@ private static VersionInfo GetVersionInfo()
             loaderAssemblyName: typeof(OnnxTransform).Assembly.FullName);
         }
 
-        public static IDataTransform Create(IHostEnvironment env, IDataView input, string modelFile)
+        public static IDataTransform Create(IHostEnvironment env, IDataView input, string modelFile, int? gpuDeviceId = null, bool fallbackToCpu = false)
         {
-            var args = new Arguments { ModelFile = modelFile, InputColumns = new string[] { }, OutputColumns = new string[] { } };
+            var args = new Arguments {
+                ModelFile = modelFile,
+                InputColumns = new string[] { },
+                OutputColumns = new string[] { },
+                GpuDeviceId = gpuDeviceId,
+                FallbackToCpu = fallbackToCpu };
+
             return Create(env, args, input);
         }
 
-        public static IDataTransform Create(IHostEnvironment env, IDataView input, string modelFile, string[] inputColumns, string[] outputColumns)
+        public static IDataTransform Create(IHostEnvironment env, IDataView input, string modelFile, string[] inputColumns, string[] outputColumns, int? gpuDeviceId = null, bool fallbackToCpu = false)
         {
-            var args = new Arguments { ModelFile = modelFile, InputColumns = inputColumns, OutputColumns = outputColumns };
+            var args = new Arguments {
+                ModelFile = modelFile,
+                InputColumns = inputColumns,
+                OutputColumns = outputColumns,
+                GpuDeviceId = gpuDeviceId,
+                FallbackToCpu = fallbackToCpu };
+
             return Create(env, args, input);
         }
 
@@ -156,14 +180,21 @@ private OnnxTransform(IHostEnvironment env, Arguments args, byte[] modelBytes =
             foreach (var col in args.OutputColumns)
                 Host.CheckNonWhiteSpace(col, nameof(args.OutputColumns));
 
-            if (modelBytes == null)
+            try
+            {
+                if (modelBytes == null)
+                {
+                    Host.CheckNonWhiteSpace(args.ModelFile, nameof(args.ModelFile));
+                    Host.CheckUserArg(File.Exists(args.ModelFile), nameof(args.ModelFile));
+                    Model = new OnnxModel(args.ModelFile, args.GpuDeviceId, args.FallbackToCpu);
+                }
+                else
+                    Model = OnnxModel.CreateFromBytes(modelBytes, args.GpuDeviceId, args.FallbackToCpu);
+            }
+            catch (OnnxRuntimeException e)
             {
-                Host.CheckNonWhiteSpace(args.ModelFile, nameof(args.ModelFile));
-                Host.CheckUserArg(File.Exists(args.ModelFile), nameof(args.ModelFile));
-                Model = new OnnxModel(args.ModelFile);
+                 throw Host.Except(e, $"Error initializing model :{e.ToString()}");
             }
-            else
-                Model = OnnxModel.CreateFromBytes(modelBytes);
 
             var modelInfo = Model.ModelInfo;
             Inputs = (args.InputColumns.Count() == 0 ) ? Model.InputNames.ToArray() : args.InputColumns;
@@ -184,18 +215,68 @@ private OnnxTransform(IHostEnvironment env, Arguments args, byte[] modelBytes =
             _args = args;
         }
 
-        public OnnxTransform(IHostEnvironment env, string modelFile)
-            : this(env, new Arguments() { ModelFile = modelFile, InputColumns = new string[] { }, OutputColumns = new string[] { } })
+        /// <summary>
+        /// Transform for scoring ONNX models. Input data column names/types must exactly match
+        /// all model input names. All possible output columns are generated, with names/types
+        /// specified by model.
+        /// </summary>
+        /// <param name="env">The environment to use.</param>
+        /// <param name="modelFile">Model file path.</param>
+        /// <param name="gpuDeviceId">Optional GPU device ID to run execution on. Null for CPU.</param>
+        /// <param name="fallbackToCpu">If GPU error, raise exception or fallback to CPU.</param>
+        public OnnxTransform(IHostEnvironment env, string modelFile, int? gpuDeviceId = null, bool fallbackToCpu = false)
+            : this(env, new Arguments()
+            {
+                ModelFile = modelFile,
+                InputColumns = new string[] {},
+                OutputColumns = new string[] {},
+                GpuDeviceId = gpuDeviceId,
+                FallbackToCpu = fallbackToCpu
+            })
         {
         }
 
-        public OnnxTransform(IHostEnvironment env, string modelFile, string inputColumn, string outputColumn)
-            : this(env, new Arguments() { ModelFile = modelFile, InputColumns = new[] { inputColumn }, OutputColumns = new[] { outputColumn } })
+        /// <summary>
+        /// Transform for scoring ONNX models. Input data column name/type must exactly match
+        /// the model specification. Only 1 output column is generated.
+        /// </summary>
+        /// <param name="env">The environment to use.</param>
+        /// <param name="modelFile">Model file path.</param>
+        /// <param name="inputColumn">The name of the input data column. Must match model input name.</param>
+        /// <param name="outputColumn">The output columns to generate. Names must match model specifications. Data types are inferred from model.</param>
+        /// <param name="gpuDeviceId">Optional GPU device ID to run execution on. Null for CPU.</param>
+        /// <param name="fallbackToCpu">If GPU error, raise exception or fallback to CPU.</param>
+        public OnnxTransform(IHostEnvironment env, string modelFile, string inputColumn, string outputColumn, int? gpuDeviceId = null, bool fallbackToCpu = false)
+            : this(env, new Arguments()
+            {
+                ModelFile = modelFile,
+                InputColumns = new[] { inputColumn },
+                OutputColumns = new[] { outputColumn },
+                GpuDeviceId = gpuDeviceId,
+                FallbackToCpu = fallbackToCpu
+            })
         {
         }
 
-        public OnnxTransform(IHostEnvironment env, string modelFile, string[] inputColumns, string[] outputColumns)
-            : this(env, new Arguments() { ModelFile = modelFile, InputColumns = inputColumns, OutputColumns = outputColumns })
+        /// <summary>
+        /// Transform for scoring ONNX models. Input data column names/types must exactly match
+        /// all model input names. Only the output columns specified will be generated.
+        /// </summary>
+        /// <param name="env">The environment to use.</param>
+        /// <param name="modelFile">Model file path.</param>
+        /// <param name="inputColumns">The name of the input data columns. Must match model's input names.</param>
+        /// <param name="outputColumns">The output columns to generate. Names must match model specifications. Data types are inferred from model.</param>
+        /// <param name="gpuDeviceId">Optional GPU device ID to run execution on. Null for CPU.</param>
+        /// <param name="fallbackToCpu">If GPU error, raise exception or fallback to CPU.</param>
+        public OnnxTransform(IHostEnvironment env, string modelFile, string[] inputColumns, string[] outputColumns, int? gpuDeviceId = null, bool fallbackToCpu = false)
+            : this(env, new Arguments()
+            {
+                ModelFile = modelFile,
+                InputColumns = inputColumns,
+                OutputColumns = outputColumns,
+                GpuDeviceId = gpuDeviceId,
+                FallbackToCpu = fallbackToCpu
+            })
         {
         }
 

src/Microsoft.ML.OnnxTransform/OnnxUtils.cs

@@ -71,10 +71,36 @@ public OnnxNodeInfo(string name, OnnxShape shape, System.Type type)
         public readonly List<string> InputNames;
         public readonly List<string> OutputNames;
 
-        public OnnxModel(string modelFile)
+        /// <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)
         {
             _modelFile = modelFile;
-            _session = new InferenceSession(modelFile);
+
+            if (gpuDeviceId.HasValue)
+            {
+                try
+                {
+                    _session = new InferenceSession(modelFile, SessionOptions.MakeSessionOptionWithCudaProvider(gpuDeviceId.Value));
+                }
+                catch (OnnxRuntimeException)
+                {
+                    if (fallbackToCpu)
+                        _session = new InferenceSession(modelFile);
+                    else
+                        // if called from OnnxTranform, is caught and rethrown.
+                        throw;
+                }
+            }
+            else
+            {
+                _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();
@@ -83,16 +109,28 @@ public OnnxModel(string modelFile)
         /// <summary>
         /// Create an OnnxModel from a byte[]
         /// </summary>
-        /// <param name="modelBytes"></param>
+        /// <param name="modelBytes">Bytes of the serialized model</param>
         /// <returns>OnnxModel</returns>
         public static OnnxModel CreateFromBytes(byte[] modelBytes)
+        {
+            return CreateFromBytes(modelBytes, null, false);
+        }
+
+        /// <summary>
+        /// Create an OnnxModel from a byte[]. Set execution to GPU if required.
+        /// </summary>
+        /// <param name="modelBytes">Bytes of the serialized model.</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>
+        /// <returns>OnnxModel</returns>
+        public static OnnxModel CreateFromBytes(byte[] modelBytes, int? gpuDeviceId = null, bool fallbackToCpu = false)
         {
             var tempModelDir = Path.Combine(Path.GetTempPath(), Guid.NewGuid().ToString());
             Directory.CreateDirectory(tempModelDir);
 
             var tempModelFile = Path.Combine(tempModelDir, "model.onnx");
             File.WriteAllBytes(tempModelFile, modelBytes);
-            return new OnnxModel(tempModelFile);
+            return new OnnxModel(tempModelFile, gpuDeviceId, fallbackToCpu);
 
             // TODO:
             // tempModelFile is needed in case the model needs to be saved
@@ -103,8 +141,8 @@ public static OnnxModel CreateFromBytes(byte[] modelBytes)
         /// <summary>
         /// Uses an open session to score a list of NamedOnnxValues.
         /// </summary>
-        /// <param name="inputNamedOnnxValues">The NamedOnnxValues to score</param>
-        /// <returns>Resulting output NamedOnnxValues list</returns>
+        /// <param name="inputNamedOnnxValues">The NamedOnnxValues to score.</param>
+        /// <returns>Resulting output NamedOnnxValues list.</returns>
         public IReadOnlyCollection<NamedOnnxValue> Run(List<NamedOnnxValue> inputNamedOnnxValues)
         {
             return _session.Run(inputNamedOnnxValues);
@@ -170,9 +208,9 @@ internal sealed class OnnxUtils
         /// <summary>
         /// Creates a NamedOnnxValue from a scalar value.
         /// </summary>
-        /// <typeparam name="T">The type of the Tensor contained in the NamedOnnxValue</typeparam>
-        /// <param name="name">The name of the NamedOnnxValue</param>
-        /// <param name="data">The data values of the Tensor</param>
+        /// <typeparam name="T">The type of the Tensor contained in the NamedOnnxValue.</typeparam>
+        /// <param name="name">The name of the NamedOnnxValue.</param>
+        /// <param name="data">The data values of the Tensor.</param>
         /// <returns>NamedOnnxValue</returns>
         public static NamedOnnxValue CreateScalarNamedOnnxValue<T>(string name, T data)
         {
@@ -185,10 +223,10 @@ public static NamedOnnxValue CreateScalarNamedOnnxValue<T>(string name, T data)
         /// Create a NamedOnnxValue from vbuffer span. Checks if the tensor type
         /// is supported by OnnxRuntime prior to execution.
         /// </summary>
-        /// <typeparam name="T">The type of the Tensor contained in the NamedOnnxValue</typeparam>
-        /// <param name="name">The name of the NamedOnnxValue</param>
+        /// <typeparam name="T">The type of the Tensor contained in the NamedOnnxValue.</typeparam>
+        /// <param name="name">The name of the NamedOnnxValue.</param>
         /// <param name="data">A span containing the data</param>
-        /// <param name="shape">The shape of the Tensor being created</param>
+        /// <param name="shape">The shape of the Tensor being created.</param>
         /// <returns>NamedOnnxValue</returns>
         public static NamedOnnxValue CreateNamedOnnxValue<T>(string name, ReadOnlySpan<T> data, OnnxShape shape)
         {

test/Microsoft.ML.OnnxTransformTest/OnnxTransformTests.cs

@@ -236,7 +236,7 @@ void TestCommandLine()
                 return;
 
             var env = new MLContext();
-            var x = Maml.Main(new[] { @"showschema loader=Text{col=data_0:R4:0-150527} xf=Onnx{InputColumns={data_0} OutputColumns={softmaxout_1} model={squeezenet/00000001/model.onnx}}" });
+            var x = Maml.Main(new[] { @"showschema loader=Text{col=data_0:R4:0-150527} xf=Onnx{InputColumns={data_0} OutputColumns={softmaxout_1} model={squeezenet/00000001/model.onnx} GpuDeviceId=0 FallbackToCpu=+}" });
             Assert.Equal(0, x);
         }