@@ -23,6 +23,11 @@ namespace Microsoft.ML.Tests
2323{
2424 public class OnnxTransformTests : TestDataPipeBase
2525 {
26+ // These two members are meant to be changed
27+ // Only when manually testing the Onnx GPU nuggets
28+ private const bool _fallbackToCpu = true ;
29+ private static int ? _gpuDeviceId = null ;
30+
2631 private const int InputSize = 150528 ;
2732
2833 private class TestData
@@ -134,7 +139,7 @@ public void TestSimpleCase()
134139 var xyData = new List < TestDataXY > { new TestDataXY ( ) { A = new float [ InputSize ] } } ;
135140 var stringData = new List < TestDataDifferntType > { new TestDataDifferntType ( ) { data_0 = new string [ InputSize ] } } ;
136141 var sizeData = new List < TestDataSize > { new TestDataSize ( ) { data_0 = new float [ 2 ] } } ;
137- var pipe = ML . Transforms . ApplyOnnxModel ( new [ ] { "softmaxout_1" } , new [ ] { "data_0" } , modelFile ) ;
142+ var pipe = ML . Transforms . ApplyOnnxModel ( new [ ] { "softmaxout_1" } , new [ ] { "data_0" } , modelFile , gpuDeviceId : _gpuDeviceId , fallbackToCpu : _fallbackToCpu ) ;
138143
139144 var invalidDataWrongNames = ML . Data . LoadFromEnumerable ( xyData ) ;
140145 var invalidDataWrongTypes = ML . Data . LoadFromEnumerable ( stringData ) ;
@@ -231,7 +236,7 @@ public void OnnxWorkout()
231236 var pipe = ML . Transforms . LoadImages ( "data_0" , imageFolder , "imagePath" )
232237 . Append ( ML . Transforms . ResizeImages ( "data_0" , imageHeight , imageWidth ) )
233238 . Append ( ML . Transforms . ExtractPixels ( "data_0" , interleavePixelColors : true ) )
234- . Append ( ML . Transforms . ApplyOnnxModel ( "softmaxout_1" , "data_0" , modelFile ) ) ;
239+ . Append ( ML . Transforms . ApplyOnnxModel ( "softmaxout_1" , "data_0" , modelFile , gpuDeviceId : _gpuDeviceId , fallbackToCpu : _fallbackToCpu ) ) ;
235240
236241 TestEstimatorCore ( pipe , data ) ;
237242
@@ -292,7 +297,7 @@ public void OnnxModelScenario()
292297 }
293298 } ) ;
294299
295- var pipeline = ML . Transforms . ApplyOnnxModel ( "softmaxout_1" , "data_0" , modelFile ) ;
300+ var pipeline = ML . Transforms . ApplyOnnxModel ( "softmaxout_1" , "data_0" , modelFile , gpuDeviceId : _gpuDeviceId , fallbackToCpu : _fallbackToCpu ) ;
296301 var onnxTransformer = pipeline . Fit ( dataView ) ;
297302 var onnx = onnxTransformer . Transform ( dataView ) ;
298303 var scoreCol = onnx . Schema [ "softmaxout_1" ] ;
@@ -325,7 +330,7 @@ public void OnnxModelMultiInput()
325330 inb = new float [ ] { 1 , 2 , 3 , 4 , 5 }
326331 }
327332 } ) ;
328- var pipeline = ML . Transforms . ApplyOnnxModel ( new [ ] { "outa" , "outb" } , new [ ] { "ina" , "inb" } , modelFile ) ;
333+ var pipeline = ML . Transforms . ApplyOnnxModel ( new [ ] { "outa" , "outb" } , new [ ] { "ina" , "inb" } , modelFile , gpuDeviceId : _gpuDeviceId , fallbackToCpu : _fallbackToCpu ) ;
329334 var onnxTransformer = pipeline . Fit ( dataView ) ;
330335 var onnx = onnxTransformer . Transform ( dataView ) ;
331336
@@ -365,7 +370,7 @@ public void OnnxModelOutputDifferentOrder()
365370 }
366371 } ) ;
367372 // The model returns the output columns in the order outa, outb. We are doing the opposite here, making sure the name mapping is correct.
368- var pipeline = ML . Transforms . ApplyOnnxModel ( new [ ] { "outb" , "outa" } , new [ ] { "ina" , "inb" } , modelFile ) ;
373+ var pipeline = ML . Transforms . ApplyOnnxModel ( new [ ] { "outb" , "outa" } , new [ ] { "ina" , "inb" } , modelFile , gpuDeviceId : _gpuDeviceId , fallbackToCpu : _fallbackToCpu ) ;
369374 var onnxTransformer = pipeline . Fit ( dataView ) ;
370375 var onnx = onnxTransformer . Transform ( dataView ) ;
371376
@@ -391,7 +396,7 @@ public void OnnxModelOutputDifferentOrder()
391396 ( onnxTransformer as IDisposable ) ? . Dispose ( ) ;
392397
393398 // The model returns the output columns in the order outa, outb. We are doing only a subset, outb, to make sure the mapping works.
394- pipeline = ML . Transforms . ApplyOnnxModel ( new [ ] { "outb" } , new [ ] { "ina" , "inb" } , modelFile ) ;
399+ pipeline = ML . Transforms . ApplyOnnxModel ( new [ ] { "outb" } , new [ ] { "ina" , "inb" } , modelFile , gpuDeviceId : _gpuDeviceId , fallbackToCpu : _fallbackToCpu ) ;
395400 onnxTransformer = pipeline . Fit ( dataView ) ;
396401 onnx = onnxTransformer . Transform ( dataView ) ;
397402
@@ -425,7 +430,7 @@ public void TestUnknownDimensions()
425430 new TestDataUnknownDimensions ( ) { input = new float [ ] { - 1.1f , - 1.3f , 1.2f } } ,
426431 } ;
427432 var idv = mlContext . Data . LoadFromEnumerable ( data ) ;
428- var pipeline = ML . Transforms . ApplyOnnxModel ( modelFile ) ;
433+ var pipeline = ML . Transforms . ApplyOnnxModel ( modelFile , gpuDeviceId : _gpuDeviceId , fallbackToCpu : _fallbackToCpu ) ;
429434 var onnxTransformer = pipeline . Fit ( idv ) ;
430435 var transformedValues = onnxTransformer . Transform ( idv ) ;
431436 var predictions = mlContext . Data . CreateEnumerable < PredictionUnknownDimensions > ( transformedValues , reuseRowObject : false ) . ToArray ( ) ;
@@ -451,7 +456,7 @@ public void TestOnnxNoneDimValue()
451456 new TestDataNoneDimension ( ) { features = new float [ ] { 6.3f , 3.3f , 6.0f , 2.5f } } ,
452457 } ;
453458 var idv = mlContext . Data . LoadFromEnumerable ( data ) ;
454- var pipeline = ML . Transforms . ApplyOnnxModel ( modelFile ) ;
459+ var pipeline = ML . Transforms . ApplyOnnxModel ( modelFile , gpuDeviceId : _gpuDeviceId , fallbackToCpu : _fallbackToCpu ) ;
455460 var onnxTransformer = pipeline . Fit ( idv ) ;
456461 var transformedValues = onnxTransformer . Transform ( idv ) ;
457462 var predictions = mlContext . Data . CreateEnumerable < PredictionNoneDimension > ( transformedValues , reuseRowObject : false ) . ToArray ( ) ;
@@ -526,7 +531,7 @@ public void OnnxModelInMemoryImage()
526531 // "softmaxout_1" are model input and output names stored in the used ONNX model file. Users may need to inspect their own models to
527532 // get the right input and output column names.
528533 var pipeline = ML . Transforms . ExtractPixels ( "data_0" , "Image" ) // Map column "Image" to column "data_0"
529- . Append ( ML . Transforms . ApplyOnnxModel ( "softmaxout_1" , "data_0" , modelFile ) ) ; // Map column "data_0" to column "softmaxout_1"
534+ . Append ( ML . Transforms . ApplyOnnxModel ( "softmaxout_1" , "data_0" , modelFile , gpuDeviceId : _gpuDeviceId , fallbackToCpu : _fallbackToCpu ) ) ; // Map column "data_0" to column "softmaxout_1"
530535 var model = pipeline . Fit ( dataView ) ;
531536 var onnx = model . Transform ( dataView ) ;
532537
@@ -576,7 +581,7 @@ public void TestOnnxZipMapWithInt64Keys()
576581 } ;
577582
578583 var dataView = ML . Data . LoadFromEnumerable ( dataPoints ) ;
579- var pipeline = ML . Transforms . ApplyOnnxModel ( new [ ] { "output" } , new [ ] { "input" } , modelFile ) ;
584+ var pipeline = ML . Transforms . ApplyOnnxModel ( new [ ] { "output" } , new [ ] { "input" } , modelFile , gpuDeviceId : _gpuDeviceId , fallbackToCpu : _fallbackToCpu ) ;
580585 var onnxTransformer = pipeline . Fit ( dataView ) ;
581586 var transformedDataView = onnxTransformer . Transform ( dataView ) ;
582587
@@ -629,7 +634,7 @@ public void TestOnnxZipMapWithStringKeys()
629634 } ;
630635
631636 var dataView = ML . Data . LoadFromEnumerable ( dataPoints ) ;
632- var pipeline = ML . Transforms . ApplyOnnxModel ( new [ ] { "output" } , new [ ] { "input" } , modelFile ) ;
637+ var pipeline = ML . Transforms . ApplyOnnxModel ( new [ ] { "output" } , new [ ] { "input" } , modelFile , gpuDeviceId : _gpuDeviceId , fallbackToCpu : _fallbackToCpu ) ;
633638 var onnxTransformer = pipeline . Fit ( dataView ) ;
634639 var transformedDataView = onnxTransformer . Transform ( dataView ) ;
635640
@@ -794,19 +799,19 @@ public void TestOnnxTransformWithCustomShapes()
794799 // Test 1.
795800 pipeline [ 0 ] = ML . Transforms . ApplyOnnxModel (
796801 new [ ] { nameof ( PredictionWithCustomShape . argmax ) } , new [ ] { nameof ( InputWithCustomShape . input ) } ,
797- modelFile , shapeDictionary ) ;
802+ modelFile , shapeDictionary , gpuDeviceId : _gpuDeviceId , fallbackToCpu : _fallbackToCpu ) ;
798803 onnxTransformer [ 0 ] = pipeline [ 0 ] . Fit ( dataView ) ;
799804 transformedDataViews [ 0 ] = onnxTransformer [ 0 ] . Transform ( dataView ) ;
800805
801806 // Test 2.
802807 pipeline [ 1 ] = ML . Transforms . ApplyOnnxModel (
803808 nameof ( PredictionWithCustomShape . argmax ) , nameof ( InputWithCustomShape . input ) ,
804- modelFile , shapeDictionary ) ;
809+ modelFile , shapeDictionary , gpuDeviceId : _gpuDeviceId , fallbackToCpu : _fallbackToCpu ) ;
805810 onnxTransformer [ 1 ] = pipeline [ 1 ] . Fit ( dataView ) ;
806811 transformedDataViews [ 1 ] = onnxTransformer [ 1 ] . Transform ( dataView ) ;
807812
808813 // Test 3.
809- pipeline [ 2 ] = ML . Transforms . ApplyOnnxModel ( modelFile , shapeDictionary ) ;
814+ pipeline [ 2 ] = ML . Transforms . ApplyOnnxModel ( modelFile , shapeDictionary , gpuDeviceId : _gpuDeviceId , fallbackToCpu : _fallbackToCpu ) ;
810815 onnxTransformer [ 2 ] = pipeline [ 2 ] . Fit ( dataView ) ;
811816 transformedDataViews [ 2 ] = onnxTransformer [ 2 ] . Transform ( dataView ) ;
812817
@@ -856,7 +861,7 @@ private void TryModelWithCustomShapesHelper(IDictionary<string, int[]> shapeDict
856861
857862 // Define a ONNX transform, trains it, and apply it to the input data.
858863 var pipeline = ML . Transforms . ApplyOnnxModel ( new [ ] { "outa" , "outb" } , new [ ] { "ina" , "inb" } ,
859- modelFile , shapeDictionary ) ;
864+ modelFile , shapeDictionary , gpuDeviceId : _gpuDeviceId , fallbackToCpu : _fallbackToCpu ) ;
860865 }
861866
862867 /// <summary>
@@ -956,7 +961,7 @@ public void TestOnnxTransformSaveAndLoadWithCustomShapes()
956961 var dataView = ML . Data . LoadFromEnumerable ( dataPoints ) ;
957962
958963 var pipeline = ML . Transforms . ApplyOnnxModel ( nameof ( PredictionWithCustomShape . argmax ) ,
959- nameof ( InputWithCustomShape . input ) , modelFile , shapeDictionary ) ;
964+ nameof ( InputWithCustomShape . input ) , modelFile , shapeDictionary , gpuDeviceId : _gpuDeviceId , fallbackToCpu : _fallbackToCpu ) ;
960965
961966 var model = pipeline . Fit ( dataView ) ;
962967
0 commit comments