diff --git a/src/Microsoft.ML.OnnxConverter/OnnxExportExtensions.cs b/src/Microsoft.ML.OnnxConverter/OnnxExportExtensions.cs index df377e46d0..88194c35b8 100644 --- a/src/Microsoft.ML.OnnxConverter/OnnxExportExtensions.cs +++ b/src/Microsoft.ML.OnnxConverter/OnnxExportExtensions.cs @@ -47,27 +47,51 @@ private static ModelProto ConvertToOnnxProtobufCore(IHostEnvironment env, OnnxCo /// List of output columns we want to keep. /// An ONNX model equivalent to the converted ML.NET model. [BestFriend] - internal static ModelProto ConvertToOnnxProtobuf(this ModelOperationsCatalog catalog, ITransformer transform, IDataView inputData, string[] outputColumns = null) + internal static ModelProto ConvertToOnnxProtobuf(this ModelOperationsCatalog catalog, ITransformer transform, IDataView inputData, string[] outputColumns = null) => + ConvertToOnnxProtobuf(catalog, transform, inputData.Schema, outputColumns); + + /// + /// Convert the specified to ONNX format. Note that ONNX uses Google's Protobuf so the returned value is a Protobuf object. + /// + /// The class that attached to. + /// The that will be converted into ONNX format. + /// The input of the specified transform. + /// The OpSet version to use for exporting the model. This value must be greater than or equal to 9 and less than or equal to 12 + /// An ONNX model equivalent to the converted ML.NET model. + [BestFriend] + internal static ModelProto ConvertToOnnxProtobuf(this ModelOperationsCatalog catalog, ITransformer transform, IDataView inputData, int opSetVersion) => + ConvertToOnnxProtobuf(catalog, transform, inputData.Schema, opSetVersion); + + /// + /// Convert the specified to ONNX format. Note that ONNX uses Google's Protobuf so the returned value is a Protobuf object. + /// + /// The class that attached to. + /// The that will be converted into ONNX format. + /// The schema of the input to the transformer. + /// List of output columns we want to keep. + /// An ONNX model equivalent to the converted ML.NET model. + [BestFriend] + internal static ModelProto ConvertToOnnxProtobuf(this ModelOperationsCatalog catalog, ITransformer transform, DataViewSchema inputSchema, string[] outputColumns = null) { var env = catalog.GetEnvironment(); var ctx = new OnnxContextImpl(env, "model", "ML.NET", "0", 0, "machinelearning.dotnet", OnnxVersion.Stable); - return ConvertToOnnxProtobufCore(env, ctx, transform, inputData, outputColumns); + return ConvertToOnnxProtobufCore(env, ctx, transform, new EmptyDataView(env, inputSchema), outputColumns); } /// /// Convert the specified to ONNX format. Note that ONNX uses Google's Protobuf so the returned value is a Protobuf object. /// - /// The class that attached to. + /// The class that attached to. /// The that will be converted into ONNX format. - /// The input of the specified transform. + /// The schema of the input to the transformer. /// The OpSet version to use for exporting the model. This value must be greater than or equal to 9 and less than or equal to 12 /// An ONNX model equivalent to the converted ML.NET model. [BestFriend] - internal static ModelProto ConvertToOnnxProtobuf(this ModelOperationsCatalog catalog, ITransformer transform, IDataView inputData, int opSetVersion) + internal static ModelProto ConvertToOnnxProtobuf(this ModelOperationsCatalog catalog, ITransformer transform, DataViewSchema inputSchema, int opSetVersion) { var env = catalog.GetEnvironment(); var ctx = new OnnxContextImpl(env, "model", "ML.NET", "0", 0, "machinelearning.dotnet", OnnxVersion.Stable, opSetVersion); - return ConvertToOnnxProtobufCore(env, ctx, transform, inputData); + return ConvertToOnnxProtobufCore(env, ctx, transform, new EmptyDataView(env, inputSchema)); } /// @@ -104,5 +128,40 @@ public static void ConvertToOnnx(this ModelOperationsCatalog catalog, ITransform /// An ONNX model equivalent to the converted ML.NET model. public static void ConvertToOnnx(this ModelOperationsCatalog catalog, ITransformer transform, IDataView inputData, Stream stream, params string[] outputColumns) => ConvertToOnnxProtobuf(catalog, transform, inputData, outputColumns).WriteTo(stream); + + /// + /// Convert the specified to ONNX format and writes to a stream. + /// + /// The class that attached to. + /// The that will be converted into ONNX format. + /// The schema of the input to the transformer. + /// The stream to write the protobuf model to. + /// An ONNX model equivalent to the converted ML.NET model. + public static void ConvertToOnnx(this ModelOperationsCatalog catalog, ITransformer transform, DataViewSchema inputSchema, Stream stream) => + ConvertToOnnxProtobuf(catalog, transform, inputSchema).WriteTo(stream); + + /// + /// Convert the specified to ONNX format and writes to a stream. + /// + /// The class that attached to. + /// The that will be converted into ONNX format. + /// The schema of the input to the transformer. + /// The OpSet version to use for exporting the model. This value must be greater than or equal to 9 and less than or equal to 12 + /// The stream to write the protobuf model to. + /// An ONNX model equivalent to the converted ML.NET model. + public static void ConvertToOnnx(this ModelOperationsCatalog catalog, ITransformer transform, DataViewSchema inputSchema, int opSetVersion, Stream stream) => + ConvertToOnnxProtobuf(catalog, transform, inputSchema, opSetVersion).WriteTo(stream); + + /// + /// Convert the specified to ONNX format and writes to a stream. + /// + /// The class that attached to. + /// The that will be converted into ONNX format. + /// The schema of the input to the transformer. + /// The stream to write the protobuf model to. + /// List of output columns we want to keep. + /// An ONNX model equivalent to the converted ML.NET model. + public static void ConvertToOnnx(this ModelOperationsCatalog catalog, ITransformer transform, DataViewSchema inputSchema, Stream stream, params string[] outputColumns) => + ConvertToOnnxProtobuf(catalog, transform, inputSchema, outputColumns).WriteTo(stream); } }