diff --git a/src/Microsoft.ML.OnnxTransformer/OnnxCatalog.cs b/src/Microsoft.ML.OnnxTransformer/OnnxCatalog.cs
index 39bc1fc05d..a1c27b2e61 100644
--- a/src/Microsoft.ML.OnnxTransformer/OnnxCatalog.cs
+++ b/src/Microsoft.ML.OnnxTransformer/OnnxCatalog.cs
@@ -96,8 +96,19 @@ public static OnnxScoringEstimator ApplyOnnxModel(this TransformsCatalog catalog
string modelFile,
int? gpuDeviceId = null,
bool fallbackToCpu = false)
- => new OnnxScoringEstimator(CatalogUtils.GetEnvironment(catalog), new[] { outputColumnName }, new[] { inputColumnName },
- modelFile, gpuDeviceId, fallbackToCpu);
+ => new OnnxScoringEstimator(CatalogUtils.GetEnvironment(catalog), new[] { outputColumnName }, new[] { inputColumnName },
+ modelFile, gpuDeviceId, fallbackToCpu);
+
+ ///
+ /// Create a using the specified .
+ /// Please refer to to learn more about the necessary dependencies,
+ /// and how to run it on a GPU.
+ ///
+ /// The transform's catalog.
+ /// Options for the .
+ public static OnnxScoringEstimator ApplyOnnxModel(this TransformsCatalog catalog, OnnxOptions options)
+ => new OnnxScoringEstimator(CatalogUtils.GetEnvironment(catalog), options.OutputColumns, options.InputColumns, options.ModelFile,
+ options.GpuDeviceId, options.FallbackToCpu, options.ShapeDictionary, options.RecursionLimit, options.InterOpNumThreads, options.IntraOpNumThreads);
///
/// Create a , which applies a pre-trained Onnx model to the column.
diff --git a/src/Microsoft.ML.OnnxTransformer/OnnxOptions.cs b/src/Microsoft.ML.OnnxTransformer/OnnxOptions.cs
new file mode 100644
index 0000000000..ad04ff23b6
--- /dev/null
+++ b/src/Microsoft.ML.OnnxTransformer/OnnxOptions.cs
@@ -0,0 +1,59 @@
+// 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;
+
+namespace Microsoft.ML.Transforms.Onnx
+{
+ ///
+ /// The options for an .
+ ///
+ public sealed class OnnxOptions
+ {
+ ///
+ /// Path to the onnx model file.
+ ///
+ public string ModelFile;
+
+ ///
+ /// Name of the input column.
+ ///
+ public string[] InputColumns;
+
+ ///
+ /// Name of the output column.
+ ///
+ public string[] OutputColumns;
+
+ ///
+ /// GPU device id to run on (e.g. 0,1,..). Null for CPU. Requires CUDA 9.1.
+ ///
+ public int? GpuDeviceId = null;
+
+ ///
+ /// If true, resumes execution on CPU upon GPU error. If false, will raise the GPU exception.
+ ///
+ public bool FallbackToCpu = false;
+
+ ///
+ /// ONNX shapes to be used over those loaded from .
+ ///
+ public IDictionary ShapeDictionary;
+
+ ///
+ /// Protobuf CodedInputStream recursion limit.
+ ///
+ public int RecursionLimit = 100;
+
+ ///
+ /// Controls the number of threads used to parallelize the execution of the graph (across nodes).
+ ///
+ public int? InterOpNumThreads = null;
+
+ ///
+ /// Controls the number of threads to use to run the model.
+ ///
+ public int? IntraOpNumThreads = null;
+ }
+}
diff --git a/src/Microsoft.ML.OnnxTransformer/OnnxTransform.cs b/src/Microsoft.ML.OnnxTransformer/OnnxTransform.cs
index bed90e9fec..e8be80536b 100644
--- a/src/Microsoft.ML.OnnxTransformer/OnnxTransform.cs
+++ b/src/Microsoft.ML.OnnxTransformer/OnnxTransform.cs
@@ -90,6 +90,12 @@ internal sealed class Options : TransformInputBase
[Argument(ArgumentType.AtMostOnce, HelpText = "Protobuf CodedInputStream recursion limit.", SortOrder = 6)]
public int RecursionLimit = 100;
+
+ [Argument(ArgumentType.AtMostOnce, HelpText = "Controls the number of threads used to parallelize the execution of the graph (across nodes).", SortOrder = 7)]
+ public int? InterOpNumThreads = null;
+
+ [Argument(ArgumentType.AtMostOnce, HelpText = "Controls the number of threads to use to run the model.", SortOrder = 8)]
+ public int? IntraOpNumThreads = null;
}
///
@@ -244,7 +250,8 @@ private OnnxTransformer(IHostEnvironment env, Options options, byte[] modelBytes
Host.CheckNonWhiteSpace(options.ModelFile, nameof(options.ModelFile));
Host.CheckIO(File.Exists(options.ModelFile), "Model file {0} does not exists.", options.ModelFile);
// Because we cannot delete the user file, ownModelFile should be false.
- Model = new OnnxModel(options.ModelFile, options.GpuDeviceId, options.FallbackToCpu, ownModelFile: false, shapeDictionary: shapeDictionary, options.RecursionLimit);
+ Model = new OnnxModel(options.ModelFile, options.GpuDeviceId, options.FallbackToCpu, ownModelFile: false, shapeDictionary: shapeDictionary, options.RecursionLimit,
+ options.InterOpNumThreads, options.IntraOpNumThreads);
}
else
{
@@ -309,8 +316,11 @@ internal OnnxTransformer(IHostEnvironment env, string modelFile, int? gpuDeviceI
/// If GPU error, raise exception or fallback to CPU.
///
/// Optional, specifies the Protobuf CodedInputStream recursion limit. Default value is 100.
+ /// Controls the number of threads used to parallelize the execution of the graph (across nodes).
+ /// Controls the number of threads to use to run the model.
internal OnnxTransformer(IHostEnvironment env, string[] outputColumnNames, string[] inputColumnNames, string modelFile, int? gpuDeviceId = null, bool fallbackToCpu = false,
- IDictionary shapeDictionary = null, int recursionLimit = 100)
+ IDictionary shapeDictionary = null, int recursionLimit = 100,
+ int? interOpNumThreads = null, int? intraOpNumThreads = null)
: this(env, new Options()
{
ModelFile = modelFile,
@@ -319,7 +329,9 @@ internal OnnxTransformer(IHostEnvironment env, string[] outputColumnNames, strin
GpuDeviceId = gpuDeviceId,
FallbackToCpu = fallbackToCpu,
CustomShapeInfos = shapeDictionary?.Select(pair => new CustomShapeInfo(pair.Key, pair.Value)).ToArray(),
- RecursionLimit = recursionLimit
+ RecursionLimit = recursionLimit,
+ InterOpNumThreads = interOpNumThreads,
+ IntraOpNumThreads = intraOpNumThreads
})
{
}
@@ -856,9 +868,12 @@ internal OnnxScoringEstimator(IHostEnvironment env, string modelFile, int? gpuDe
/// If GPU error, raise exception or fallback to CPU.
///
/// Optional, specifies the Protobuf CodedInputStream recursion limit. Default value is 100.
+ /// Controls the number of threads used to parallelize the execution of the graph (across nodes).
+ /// Controls the number of threads to use to run the model.
internal OnnxScoringEstimator(IHostEnvironment env, string[] outputColumnNames, string[] inputColumnNames, string modelFile,
- int? gpuDeviceId = null, bool fallbackToCpu = false, IDictionary shapeDictionary = null, int recursionLimit = 100)
- : this(env, new OnnxTransformer(env, outputColumnNames, inputColumnNames, modelFile, gpuDeviceId, fallbackToCpu, shapeDictionary, recursionLimit))
+ int? gpuDeviceId = null, bool fallbackToCpu = false, IDictionary shapeDictionary = null, int recursionLimit = 100,
+ int? interOpNumThreads = null, int? intraOpNumThreads = null)
+ : this(env, new OnnxTransformer(env, outputColumnNames, inputColumnNames, modelFile, gpuDeviceId, fallbackToCpu, shapeDictionary, recursionLimit, interOpNumThreads, intraOpNumThreads))
{
}
diff --git a/src/Microsoft.ML.OnnxTransformer/OnnxUtils.cs b/src/Microsoft.ML.OnnxTransformer/OnnxUtils.cs
index 90705ec2b7..c8c9aed970 100644
--- a/src/Microsoft.ML.OnnxTransformer/OnnxUtils.cs
+++ b/src/Microsoft.ML.OnnxTransformer/OnnxUtils.cs
@@ -165,8 +165,11 @@ public OnnxVariableInfo(string name, OnnxShape shape, Type typeInOnnxRuntime, Da
/// no longer needed.
///
/// Optional, specifies the Protobuf CodedInputStream recursion limit. Default value is 100.
+ /// Controls the number of threads used to parallelize the execution of the graph (across nodes).
+ /// Controls the number of threads to use to run the model.
public OnnxModel(string modelFile, int? gpuDeviceId = null, bool fallbackToCpu = false,
- bool ownModelFile = false, IDictionary shapeDictionary = null, int recursionLimit = 100)
+ bool ownModelFile = false, IDictionary shapeDictionary = null, int recursionLimit = 100,
+ int? interOpNumThreads = null, int? intraOpNumThreads = null)
{
// If we don't own the model file, _disposed should be false to prevent deleting user's file.
_disposed = false;
@@ -181,7 +184,14 @@ public OnnxModel(string modelFile, int? gpuDeviceId = null, bool fallbackToCpu =
catch (OnnxRuntimeException)
{
if (fallbackToCpu)
- _session = new InferenceSession(modelFile);
+ {
+ var sessionOptions = new SessionOptions()
+ {
+ InterOpNumThreads = interOpNumThreads.GetValueOrDefault(),
+ IntraOpNumThreads = intraOpNumThreads.GetValueOrDefault()
+ };
+ _session = new InferenceSession(modelFile, sessionOptions);
+ }
else
// If called from OnnxTransform, is caught and rethrown
throw;
@@ -189,7 +199,12 @@ public OnnxModel(string modelFile, int? gpuDeviceId = null, bool fallbackToCpu =
}
else
{
- _session = new InferenceSession(modelFile);
+ var sessionOptions = new SessionOptions()
+ {
+ InterOpNumThreads = interOpNumThreads.GetValueOrDefault(),
+ IntraOpNumThreads = intraOpNumThreads.GetValueOrDefault()
+ };
+ _session = new InferenceSession(modelFile, sessionOptions);
}
try
diff --git a/test/Microsoft.ML.OnnxTransformerTest/OnnxTransformTests.cs b/test/Microsoft.ML.OnnxTransformerTest/OnnxTransformTests.cs
index 48561d435d..1bc8097644 100644
--- a/test/Microsoft.ML.OnnxTransformerTest/OnnxTransformTests.cs
+++ b/test/Microsoft.ML.OnnxTransformerTest/OnnxTransformTests.cs
@@ -1,4 +1,4 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// 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.
@@ -119,8 +119,10 @@ public OnnxTransformTests(ITestOutputHelper output) : base(output)
{
}
- [OnnxFact]
- public void TestSimpleCase()
+ [OnnxTheory]
+ [InlineData(false)]
+ [InlineData(true)]
+ public void TestSimpleCase(bool useOptionsCtor)
{
var modelFile = "squeezenet/00000001/model.onnx";
var samplevector = GetSampleArrayData();
@@ -139,7 +141,19 @@ public void TestSimpleCase()
var xyData = new List { new TestDataXY() { A = new float[InputSize] } };
var stringData = new List { new TestDataDifferntType() { data_0 = new string[InputSize] } };
var sizeData = new List { new TestDataSize() { data_0 = new float[2] } };
- var pipe = ML.Transforms.ApplyOnnxModel(new[] { "softmaxout_1" }, new[] { "data_0" }, modelFile, gpuDeviceId: _gpuDeviceId, fallbackToCpu: _fallbackToCpu);
+ var options = new OnnxOptions()
+ {
+ OutputColumns = new[] { "softmaxout_1" },
+ InputColumns = new[] {"data_0" },
+ ModelFile = modelFile,
+ GpuDeviceId = _gpuDeviceId,
+ FallbackToCpu = _fallbackToCpu,
+ InterOpNumThreads = 1,
+ IntraOpNumThreads = 1
+ };
+ var pipe = useOptionsCtor ?
+ ML.Transforms.ApplyOnnxModel(options) :
+ ML.Transforms.ApplyOnnxModel(options.OutputColumns, options.InputColumns, modelFile, gpuDeviceId: _gpuDeviceId, fallbackToCpu: _fallbackToCpu);
var invalidDataWrongNames = ML.Data.LoadFromEnumerable(xyData);
var invalidDataWrongTypes = ML.Data.LoadFromEnumerable(stringData);
@@ -713,14 +727,14 @@ public void TestOnnxModelNotDisposal()
private class OnnxMapInput
{
- [OnnxMapType(typeof(int),typeof(float))]
- public IDictionary Input { get; set; }
+ [OnnxMapType(typeof(int), typeof(float))]
+ public IDictionary Input { get; set; }
}
private class OnnxMapOutput
{
- [OnnxMapType(typeof(int),typeof(float))]
- public IDictionary Output { get; set; }
+ [OnnxMapType(typeof(int), typeof(float))]
+ public IDictionary Output { get; set; }
}
///
@@ -753,10 +767,10 @@ public void SmokeInMemoryOnnxMapTypeTest()
var transformedDataView = model.Transform(dataView);
var transformedDataPoints = ML.Data.CreateEnumerable(transformedDataView, false).ToList();
- for(int i = 0; i < dataPoints.Count(); ++i)
+ for (int i = 0; i < dataPoints.Count(); ++i)
{
Assert.Equal(dataPoints[i].Input.Count(), transformedDataPoints[i].Output.Count());
- foreach(var pair in dataPoints[i].Input)
+ foreach (var pair in dataPoints[i].Input)
Assert.Equal(pair.Value, transformedDataPoints[i].Output[pair.Key + 1]);
}
}
@@ -815,7 +829,7 @@ public void TestOnnxTransformWithCustomShapes()
transformedDataViews[2] = onnxTransformer[2].Transform(dataView);
// Conduct the same check for all the 3 called public APIs.
- foreach(var transformedDataView in transformedDataViews)
+ foreach (var transformedDataView in transformedDataViews)
{
var transformedDataPoints = ML.Data.CreateEnumerable(transformedDataView, false).ToList();
@@ -901,32 +915,32 @@ public void SpecifyOnnxShapes()
Assert.False(somethingWrong);
// Case 3: this shape conflicts with output shape [1, 1, 1, 5] loaded from the model.
- shapeDictionary= new Dictionary() {
+ shapeDictionary = new Dictionary() {
{ "outb", new int[] { 5, 6 } },
};
- somethingWrong= false;
+ somethingWrong = false;
try
{
TryModelWithCustomShapesHelper(shapeDictionary);
}
catch
{
- somethingWrong= true;
+ somethingWrong = true;
}
Assert.True(somethingWrong);
// Case 4: this shape works with output shape [1, 1, 1, 5] loaded from the model.
- shapeDictionary= new Dictionary() {
+ shapeDictionary = new Dictionary() {
{ "outb", new int[] { -1, -1, -1, -1 } },
};
- somethingWrong= false;
+ somethingWrong = false;
try
{
TryModelWithCustomShapesHelper(shapeDictionary);
}
catch
{
- somethingWrong= true;
+ somethingWrong = true;
}
Assert.False(somethingWrong);
}
@@ -1024,7 +1038,7 @@ public void TestOnnxTransformSaveAndLoadWithRecursionLimit()
var pipe = ML.Transforms.LoadImages("data_0", imageFolder, "imagePath")
.Append(ML.Transforms.ResizeImages("data_0", imageHeight, imageWidth))
.Append(ML.Transforms.ExtractPixels("data_0", interleavePixelColors: true))
- .Append(ML.Transforms.ApplyOnnxModel(new []{ "softmaxout_1" }, new []{ "data_0" }, modelFile,
+ .Append(ML.Transforms.ApplyOnnxModel(new[] { "softmaxout_1" }, new[] { "data_0" }, modelFile,
gpuDeviceId: _gpuDeviceId, fallbackToCpu: _fallbackToCpu, shapeDictionary: null, recursionLimit: 50));
TestEstimatorCore(pipe, data);