Skip to content

Onnx recursion limit #5840

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Jun 24, 2021
Merged

Onnx recursion limit #5840

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
36db548
onnx-recursion-limit
Crabzmatic Jun 13, 2021
97cb474
better description for parameter
Crabzmatic Jun 13, 2021
159b507
add separate method with optional recursion limit argument
Crabzmatic Jun 17, 2021
6a83a66
increment model version and add a test for recursion limit
Crabzmatic Jun 18, 2021
2a42ab2
Update OnnxTransform.cs
michaelgsharp Jun 23, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 30 additions & 3 deletions src/Microsoft.ML.OnnxTransformer/OnnxCatalog.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ public static OnnxScoringEstimator ApplyOnnxModel(this TransformsCatalog catalog
/// <param name="modelFile">The path of the file containing the ONNX model.</param>
/// <param name="shapeDictionary">ONNX shapes to be used over those loaded from <paramref name="modelFile"/>.
/// For keys use names as stated in the ONNX model, e.g. "input". Stating the shapes with this parameter
/// is particullarly useful for working with variable dimension inputs and outputs.
/// is particularly useful for working with variable dimension inputs and outputs.
/// </param>
/// <param name="gpuDeviceId">Optional GPU device ID to run execution on, <see langword="null" /> to run on CPU.</param>
/// <param name="fallbackToCpu">If GPU error, raise exception or fallback to CPU.</param>
Expand Down Expand Up @@ -110,7 +110,7 @@ public static OnnxScoringEstimator ApplyOnnxModel(this TransformsCatalog catalog
/// <param name="modelFile">The path of the file containing the ONNX model.</param>
/// <param name="shapeDictionary">ONNX shapes to be used over those loaded from <paramref name="modelFile"/>.
/// For keys use names as stated in the ONNX model, e.g. "input". Stating the shapes with this parameter
/// is particullarly useful for working with variable dimension inputs and outputs.
/// is particularly useful for working with variable dimension inputs and outputs.
/// </param>
/// <param name="gpuDeviceId">Optional GPU device ID to run execution on, <see langword="null" /> to run on CPU.</param>
/// <param name="fallbackToCpu">If GPU error, raise exception or fallback to CPU.</param>
Expand Down Expand Up @@ -162,7 +162,7 @@ public static OnnxScoringEstimator ApplyOnnxModel(this TransformsCatalog catalog
/// <param name="modelFile">The path of the file containing the ONNX model.</param>
/// <param name="shapeDictionary">ONNX shapes to be used over those loaded from <paramref name="modelFile"/>.
/// For keys use names as stated in the ONNX model, e.g. "input". Stating the shapes with this parameter
/// is particullarly useful for working with variable dimension inputs and outputs.
/// is particularly useful for working with variable dimension inputs and outputs.
/// </param>
/// <param name="gpuDeviceId">Optional GPU device ID to run execution on, <see langword="null" /> to run on CPU.</param>
/// <param name="fallbackToCpu">If GPU error, raise exception or fallback to CPU.</param>
Expand All @@ -176,6 +176,33 @@ public static OnnxScoringEstimator ApplyOnnxModel(this TransformsCatalog catalog
=> new OnnxScoringEstimator(CatalogUtils.GetEnvironment(catalog), outputColumnNames, inputColumnNames,
modelFile, gpuDeviceId, fallbackToCpu, shapeDictionary: shapeDictionary);

/// <summary>
/// Create a <see cref="OnnxScoringEstimator"/>, which applies a pre-trained Onnx model to the <paramref name="inputColumnNames"/> columns.
/// Please refer to <see cref="OnnxScoringEstimator"/> to learn more about the necessary dependencies,
/// and how to run it on a GPU.
/// </summary>
/// <param name="catalog">The transform's catalog.</param>
/// <param name="outputColumnNames">The output columns resulting from the transformation.</param>
/// <param name="inputColumnNames">The input columns.</param>
/// <param name="modelFile">The path of the file containing the ONNX model.</param>
/// <param name="shapeDictionary">ONNX shapes to be used over those loaded from <paramref name="modelFile"/>.
/// For keys use names as stated in the ONNX model, e.g. "input". Stating the shapes with this parameter
/// is particularly useful for working with variable dimension inputs and outputs.
/// </param>
/// <param name="gpuDeviceId">Optional GPU device ID to run execution on, <see langword="null" /> to run on CPU.</param>
/// <param name="fallbackToCpu">If GPU error, raise exception or fallback to CPU.</param>
/// <param name="recursionLimit">Optional, specifies the Protobuf CodedInputStream recursion limit. Default value is 100.</param>
public static OnnxScoringEstimator ApplyOnnxModel(this TransformsCatalog catalog,
string[] outputColumnNames,
string[] inputColumnNames,
string modelFile,
IDictionary<string, int[]> shapeDictionary,
int? gpuDeviceId = null,
bool fallbackToCpu = false,
int recursionLimit = 100)
=> new OnnxScoringEstimator(CatalogUtils.GetEnvironment(catalog), outputColumnNames, inputColumnNames,
modelFile, gpuDeviceId, fallbackToCpu, shapeDictionary: shapeDictionary, recursionLimit);

/// <summary>
/// Create <see cref="DnnImageFeaturizerEstimator"/>, which applies one of the pre-trained DNN models in
/// <see cref="DnnImageModelSelector"/> to featurize an image.
Expand Down
57 changes: 44 additions & 13 deletions src/Microsoft.ML.OnnxTransformer/OnnxTransform.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,9 @@ internal sealed class Options : TransformInputBase

[Argument(ArgumentType.Multiple, HelpText = "Shapes used to overwrite shapes loaded from ONNX file.", SortOrder = 5)]
public CustomShapeInfo[] CustomShapeInfos;

[Argument(ArgumentType.AtMostOnce, HelpText = "Protobuf CodedInputStream recursion limit.", SortOrder = 6)]
public int RecursionLimit = 100;
}

/// <summary>
Expand Down Expand Up @@ -126,8 +129,9 @@ private static VersionInfo GetVersionInfo()
modelSignature: "ONNXSCOR",
// version 10001 is single input & output.
// version 10002 = multiple inputs & outputs
verWrittenCur: 0x00010002,
verReadableCur: 0x00010002,
// version 10003 = custom protobuf recursion limit
verWrittenCur: 0x00010003,
verReadableCur: 0x00010003,
verWeCanReadBack: 0x00010001,
loaderSignature: LoaderSignature,
loaderAssemblyName: typeof(OnnxTransformer).Assembly.FullName);
Expand Down Expand Up @@ -184,7 +188,26 @@ private static OnnxTransformer Create(IHostEnvironment env, ModelLoadContext ctx
}
}

var options = new Options() { InputColumns = inputs, OutputColumns = outputs, CustomShapeInfos = loadedCustomShapeInfos };
int recursionLimit;

// Recursion limit change
if (ctx.Header.ModelVerWritten >= 0x00010003)
{
recursionLimit = ctx.Reader.ReadInt32();
}
else
{
// Default if not written inside ONNX model
recursionLimit = 100;
}

var options = new Options()
{
InputColumns = inputs,
OutputColumns = outputs,
CustomShapeInfos = loadedCustomShapeInfos,
RecursionLimit = recursionLimit
};

return new OnnxTransformer(env, options, modelBytes);
}
Expand Down Expand Up @@ -221,13 +244,13 @@ 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);
Model = new OnnxModel(options.ModelFile, options.GpuDeviceId, options.FallbackToCpu, ownModelFile: false, shapeDictionary: shapeDictionary, options.RecursionLimit);
}
else
{
// Entering this region means that the byte[] is passed as the model. To feed that byte[] to ONNXRuntime, we need
// to create a temporal file to store it and then call ONNXRuntime's API to load that file.
Model = OnnxModel.CreateFromBytes(modelBytes, env, options.GpuDeviceId, options.FallbackToCpu, shapeDictionary: shapeDictionary);
Model = OnnxModel.CreateFromBytes(modelBytes, env, options.GpuDeviceId, options.FallbackToCpu, shapeDictionary: shapeDictionary, options.RecursionLimit);
}
}
catch (OnnxRuntimeException e)
Expand Down Expand Up @@ -258,16 +281,18 @@ private OnnxTransformer(IHostEnvironment env, Options options, byte[] modelBytes
/// <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>
/// <param name="shapeDictionary"></param>
/// <param name="recursionLimit">Optional, specifies the Protobuf CodedInputStream recursion limit. Default value is 100.</param>
internal OnnxTransformer(IHostEnvironment env, string modelFile, int? gpuDeviceId = null,
bool fallbackToCpu = false, IDictionary<string, int[]> shapeDictionary = null)
bool fallbackToCpu = false, IDictionary<string, int[]> shapeDictionary = null, int recursionLimit = 100)
: this(env, new Options()
{
ModelFile = modelFile,
InputColumns = new string[] { },
OutputColumns = new string[] { },
GpuDeviceId = gpuDeviceId,
FallbackToCpu = fallbackToCpu,
CustomShapeInfos = shapeDictionary?.Select(pair => new CustomShapeInfo(pair.Key, pair.Value)).ToArray()
CustomShapeInfos = shapeDictionary?.Select(pair => new CustomShapeInfo(pair.Key, pair.Value)).ToArray(),
RecursionLimit = recursionLimit
})
{
}
Expand All @@ -283,16 +308,18 @@ internal OnnxTransformer(IHostEnvironment env, string modelFile, int? gpuDeviceI
/// <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>
/// <param name="shapeDictionary"></param>
/// <param name="recursionLimit">Optional, specifies the Protobuf CodedInputStream recursion limit. Default value is 100.</param>
internal OnnxTransformer(IHostEnvironment env, string[] outputColumnNames, string[] inputColumnNames, string modelFile, int? gpuDeviceId = null, bool fallbackToCpu = false,
IDictionary<string, int[]> shapeDictionary = null)
IDictionary<string, int[]> shapeDictionary = null, int recursionLimit = 100)
: this(env, new Options()
{
ModelFile = modelFile,
InputColumns = inputColumnNames,
OutputColumns = outputColumnNames,
GpuDeviceId = gpuDeviceId,
FallbackToCpu = fallbackToCpu,
CustomShapeInfos = shapeDictionary?.Select(pair => new CustomShapeInfo(pair.Key, pair.Value)).ToArray()
CustomShapeInfos = shapeDictionary?.Select(pair => new CustomShapeInfo(pair.Key, pair.Value)).ToArray(),
RecursionLimit = recursionLimit
})
{
}
Expand Down Expand Up @@ -325,6 +352,8 @@ private protected override void SaveModel(ModelSaveContext ctx)
ctx.SaveNonEmptyString(info.Name);
ctx.Writer.WriteIntArray(info.Shape);
}

ctx.Writer.Write(_options.RecursionLimit);
}

private protected override IRowMapper MakeRowMapper(DataViewSchema inputSchema) => new Mapper(this, inputSchema);
Expand Down Expand Up @@ -807,10 +836,11 @@ public sealed class OnnxScoringEstimator : TrivialEstimator<OnnxTransformer>
/// <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>
/// <param name="shapeDictionary"></param>
/// <param name="recursionLimit">Optional, specifies the Protobuf CodedInputStream recursion limit. Default value is 100.</param>
[BestFriend]
internal OnnxScoringEstimator(IHostEnvironment env, string modelFile, int? gpuDeviceId = null, bool fallbackToCpu = false,
IDictionary<string, int[]> shapeDictionary = null)
: this(env, new OnnxTransformer(env, new string[] { }, new string[] { }, modelFile, gpuDeviceId, fallbackToCpu, shapeDictionary))
IDictionary<string, int[]> shapeDictionary = null, int recursionLimit = 100)
: this(env, new OnnxTransformer(env, new string[] { }, new string[] { }, modelFile, gpuDeviceId, fallbackToCpu, shapeDictionary, recursionLimit))
{
}

Expand All @@ -825,9 +855,10 @@ internal OnnxScoringEstimator(IHostEnvironment env, string modelFile, int? gpuDe
/// <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>
/// <param name="shapeDictionary"></param>
/// <param name="recursionLimit">Optional, specifies the Protobuf CodedInputStream recursion limit. Default value is 100.</param>
internal OnnxScoringEstimator(IHostEnvironment env, string[] outputColumnNames, string[] inputColumnNames, string modelFile,
int? gpuDeviceId = null, bool fallbackToCpu = false, IDictionary<string, int[]> shapeDictionary = null)
: this(env, new OnnxTransformer(env, outputColumnNames, inputColumnNames, modelFile, gpuDeviceId, fallbackToCpu, shapeDictionary))
int? gpuDeviceId = null, bool fallbackToCpu = false, IDictionary<string, int[]> shapeDictionary = null, int recursionLimit = 100)
: this(env, new OnnxTransformer(env, outputColumnNames, inputColumnNames, modelFile, gpuDeviceId, fallbackToCpu, shapeDictionary, recursionLimit))
{
}

Expand Down
14 changes: 8 additions & 6 deletions src/Microsoft.ML.OnnxTransformer/OnnxUtils.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -164,8 +164,9 @@ public OnnxVariableInfo(string name, OnnxShape shape, Type typeInOnnxRuntime, Da
/// <param name="ownModelFile">If true, the <paramref name="modelFile"/> will be deleted when <see cref="OnnxModel"/> is
/// no longer needed.</param>
/// <param name="shapeDictionary"></param>
/// <param name="recursionLimit">Optional, specifies the Protobuf CodedInputStream recursion limit. Default value is 100.</param>
public OnnxModel(string modelFile, int? gpuDeviceId = null, bool fallbackToCpu = false,
bool ownModelFile=false, IDictionary<string, int[]> shapeDictionary = null)
bool ownModelFile=false, IDictionary<string, int[]> shapeDictionary = null, int recursionLimit = 100)
{
// If we don't own the model file, _disposed should be false to prevent deleting user's file.
_disposed = false;
Expand Down Expand Up @@ -204,7 +205,7 @@ public OnnxModel(string modelFile, int? gpuDeviceId = null, bool fallbackToCpu =

// The CodedInputStream auto closes the stream, and we need to make sure that our main stream stays open, so creating a new one here.
using (var modelStream = new FileStream(modelFile, FileMode.Open, FileAccess.Read, FileShare.Delete | FileShare.Read))
using (var codedStream = Google.Protobuf.CodedInputStream.CreateWithLimits(modelStream, Int32.MaxValue, 100))
using (var codedStream = Google.Protobuf.CodedInputStream.CreateWithLimits(modelStream, Int32.MaxValue, recursionLimit))
model = OnnxCSharpToProtoWrapper.ModelProto.Parser.ParseFrom(codedStream);

// Parse actual input and output types stored in the loaded ONNX model to get their DataViewType's.
Expand Down Expand Up @@ -321,7 +322,7 @@ private static bool CheckOnnxShapeCompatibility(IEnumerable<int> left, IEnumerab

/// <summary>
/// Create an OnnxModel from a byte[]. Usually, a ONNX model is consumed by <see cref="OnnxModel"/> as a file.
/// With <see cref="CreateFromBytes(byte[], IHostEnvironment)"/> and <see cref="CreateFromBytes(byte[], IHostEnvironment, int?, bool, IDictionary{string, int[]})"/>,
/// With <see cref="CreateFromBytes(byte[], IHostEnvironment)"/> and <see cref="CreateFromBytes(byte[], IHostEnvironment, int?, bool, IDictionary{string, int[]}, int)"/>,
/// 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>
Expand All @@ -335,7 +336,7 @@ public static OnnxModel CreateFromBytes(byte[] modelBytes, IHostEnvironment env)
/// 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[], IHostEnvironment)"/> and
/// <see cref="CreateFromBytes(byte[], IHostEnvironment, int?, bool, IDictionary{string, int[]})"/>,
/// <see cref="CreateFromBytes(byte[], IHostEnvironment, int?, bool, IDictionary{string, int[]}, int)"/>,
/// 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>
Expand All @@ -345,14 +346,15 @@ public static OnnxModel CreateFromBytes(byte[] modelBytes, IHostEnvironment env)
/// <param name="shapeDictionary">User-provided shapes. If the key "myTensorName" is associated
/// with the value [1, 3, 5], the shape of "myTensorName" will be set to [1, 3, 5].
/// The shape loaded from <paramref name="modelBytes"/> would be overwritten.</param>
/// <param name="recursionLimit">Optional, specifies the Protobuf CodedInputStream recursion limit. Default value is 100.</param>
/// <returns>An <see cref="OnnxModel"/></returns>
public static OnnxModel CreateFromBytes(byte[] modelBytes, IHostEnvironment env, int? gpuDeviceId = null, bool fallbackToCpu = false,
IDictionary<string, int[]> shapeDictionary = null)
IDictionary<string, int[]> shapeDictionary = null, int recursionLimit = 100)
{
var tempModelFile = Path.Combine(((IHostEnvironmentInternal)env).TempFilePath, Path.GetRandomFileName());
File.WriteAllBytes(tempModelFile, modelBytes);
return new OnnxModel(tempModelFile, gpuDeviceId, fallbackToCpu,
ownModelFile: true, shapeDictionary: shapeDictionary);
ownModelFile: true, shapeDictionary: shapeDictionary, recursionLimit);
}

/// <summary>
Expand Down
Loading