Support dml quantization and add example (#1403)

mengniwang95 · web-flow · commit db0fef72137e · 2023-11-22T14:50:15.000+08:00
Signed-off-by: Mengni Wang &lt;mengni.wang@intel.com&gt;
diff --git a/docs/source/quantization.md b/docs/source/quantization.md
@@ -540,10 +540,15 @@ conf = PostTrainingQuantConfig(backend="ipex")
 conf = PostTrainingQuantConfig(backend="onnxrt_cuda_ep", device="gpu")
 ```
 ```python
+# run with ONNXRT DmlExecutionProvider on NPU
+conf = PostTrainingQuantConfig(backend="onnxrt_dml_ep", device="npu")
+```
+```python
 # run with ITEX on GPU
 conf = PostTrainingQuantConfig(backend="itex", device="gpu")
 ```
 
 ## Examples
 
 User could refer to [examples](https://github.com/intel/neural-compressor/blob/master/examples/README.md) on how to quantize a new model.
+If user wants to quantize an onnx model with npu, please refer to this [example](../../examples/onnxrt/image_recognition/onnx_model_zoo/shufflenet/quantization/ptq_static/README.md).
diff --git a/examples/onnxrt/image_recognition/onnx_model_zoo/shufflenet/quantization/ptq_static/README.md b/examples/onnxrt/image_recognition/onnx_model_zoo/shufflenet/quantization/ptq_static/README.md
@@ -11,6 +11,14 @@ pip install neural-compressor
 pip install -r requirements.txt
 ```
 
+If you want to use npu device with DmlExecutionProvider, please build onnxruntime from source on windows:
+```shell
+git clone https://github.com/microsoft/onnxruntime.git
+cd onnxruntime
+.\build.bat --use_dml --skip_tests --build_wheel --config Release
+pip install build\Windows\Release\Release\your_wheel.whl
+```
+
 > Note: Validated ONNX Runtime [Version](/docs/source/installation_guide.md#validated-software-environment).
 
 ## 2. Prepare Model
@@ -40,7 +48,8 @@ Quantize model with QLinearOps:
 bash run_quant.sh --input_model=path/to/model \  # model path as *.onnx
                    --dataset_location=/path/to/imagenet \
                    --label_path=/path/to/val.txt \
-                   --output_model=path/to/save
+                   --output_model=path/to/save \
+                   --device=cpu # default is cpu, support cpu and npu
 ```
 
 Quantize model with QDQ mode:
@@ -50,7 +59,8 @@ bash run_quant.sh --input_model=path/to/model \  # model path as *.onnx
                    --dataset_location=/path/to/imagenet \
                    --label_path=/path/to/val.txt \
                    --output_model=path/to/save \
-                   --quant_format=QDQ
+                   --quant_format=QDQ \
+                   --device=cpu # default is cpu, support cpu and npu
 ```
 
 ## 2. Benchmark
@@ -60,4 +70,8 @@ bash run_benchmark.sh --input_model=path/to/model \  # model path as *.onnx
                       --dataset_location=/path/to/imagenet \
                       --label_path=/path/to/val.txt \
                       --mode=performance # or accuracy
+                      --device=cpu # default is cpu, support cpu and npu
 ```
+
+### warning
+npu device requires the batch size of model to be 1.
diff --git a/examples/onnxrt/image_recognition/onnx_model_zoo/shufflenet/quantization/ptq_static/main.py b/examples/onnxrt/image_recognition/onnx_model_zoo/shufflenet/quantization/ptq_static/main.py
@@ -189,9 +189,10 @@ def fetcher(ids):
             except StopIteration:
                 return
 
-def eval_func(model, dataloader, metric):
+def eval_func(model, dataloader, metric, backend):
     metric.reset()
-    sess = ort.InferenceSession(model.SerializeToString(), providers=ort.get_available_providers())
+    provider = 'DmlExecutionProvider' if backend == 'onnxrt_dml_ep' else 'CPUExecutionProvider'
+    sess = ort.InferenceSession(model.SerializeToString(), providers=[provider])
     input_names = [i.name for i in sess.get_inputs()]
     for input_data, label in dataloader:
         output = sess.run(None, dict(zip(input_names, [input_data])))
@@ -247,6 +248,12 @@ def eval_func(model, dataloader, metric):
         choices=['default', 'QDQ', 'QOperator'],
         help="quantization format"
     )
+    parser.add_argument(
+        '--device',
+        type=str,
+        default='cpu',
+        choices=['cpu', 'npu'],
+    )
     parser.add_argument(
         "--batch_size",
         default=1,
@@ -257,24 +264,32 @@ def eval_func(model, dataloader, metric):
     model = onnx.load(args.model_path)
     dataloader = Dataloader(args.dataset_location, args.label_path, args.batch_size)
     top1 = TopK()
+    backend = 'default' if args.device == 'cpu' else 'onnxrt_dml_ep'
     def eval(onnx_model):
-        return eval_func(onnx_model, dataloader, top1)
+        return eval_func(onnx_model, dataloader, top1, backend)
 
     if args.benchmark:
         if args.mode == 'performance':
             from neural_compressor.benchmark import fit
             from neural_compressor.config import BenchmarkConfig
-            conf = BenchmarkConfig(warmup=10, iteration=1000, cores_per_instance=4, num_of_instance=1)
+            conf = BenchmarkConfig(warmup=10,
+                                   iteration=1000,
+                                   cores_per_instance=4,
+                                   num_of_instance=1,
+                                   device=args.device,
+                                   backend=backend)
             fit(model, conf, b_dataloader=dataloader)
         elif args.mode == 'accuracy':
             acc_result = eval(model)
             print("Batch size = %d" % dataloader.batch_size)
             print("Accuracy: %.5f" % acc_result)
     if args.tune:
         from neural_compressor import quantization, PostTrainingQuantConfig
-        config = PostTrainingQuantConfig(quant_format=args.quant_format)
+        config = PostTrainingQuantConfig(quant_format=args.quant_format,
+                                         device=args.device,
+                                         backend=backend)
  
         q_model = quantization.fit(model, config, calib_dataloader=dataloader,
 			     eval_func=eval)
 
-        q_model.save(args.output_model)
+        q_model.save(args.output_model)
diff --git a/examples/onnxrt/image_recognition/onnx_model_zoo/shufflenet/quantization/ptq_static/run_benchmark.sh b/examples/onnxrt/image_recognition/onnx_model_zoo/shufflenet/quantization/ptq_static/run_benchmark.sh
@@ -24,6 +24,9 @@ function init_params {
       --mode=*)
           mode=$(echo $var |cut -f2 -d=)
       ;;
+      --device=*)
+          device=$(echo $var |cut -f2 -d=)
+      ;;
     esac
   done
 
@@ -37,6 +40,7 @@ function run_benchmark {
             --dataset_location ${dataset_location} \
             --label_path ${label_path-${dataset_location}/../val.txt} \
             --mode=${mode} \
+            --device=${device-cpu} \
             --batch_size 1\
             --benchmark
             
diff --git a/examples/onnxrt/image_recognition/onnx_model_zoo/shufflenet/quantization/ptq_static/run_quant.sh b/examples/onnxrt/image_recognition/onnx_model_zoo/shufflenet/quantization/ptq_static/run_quant.sh
@@ -28,6 +28,9 @@ function init_params {
       --quant_format=*)
           quant_format=$(echo $var |cut -f2 -d=)
       ;;
+      --device=*)
+          device=$(echo $var |cut -f2 -d=)
+      ;;
     esac
   done
 
@@ -41,6 +44,7 @@ function run_tuning {
             --label_path ${label_path-${dataset_location}/../val.txt} \
             --output_model ${output_model} \
             --quant_format ${quant_format-default} \
+            --device ${device-cpu} \
             --tune
 }
 
diff --git a/neural_compressor/adaptor/onnxrt.py b/neural_compressor/adaptor/onnxrt.py
@@ -717,7 +717,7 @@ def _get_quantize_params(self, model, data_loader, quantize_config, iterations,
             black_nodes=black_nodes,
             white_nodes=white_nodes,
             iterations=list(range(0, iterations)),
-            backend=self.backend,
+            backend=self.backend if self.backend != "DmlExecutionProvider" else "CPUExecutionProvider",
             reduce_range=self.reduce_range,
             **kwargs,
         )

Original file line number	Diff line number	Diff line change
`@@ -717,7 +717,7 @@ def _get_quantize_params(self, model, data_loader, quantize_config, iterations,`
`717`	`717`	`black_nodes=black_nodes,`
`718`	`718`	`white_nodes=white_nodes,`
`719`	`719`	`iterations=list(range(0, iterations)),`
`720`		`- backend=self.backend,`
	`720`	`+ backend=self.backend if self.backend != "DmlExecutionProvider" else "CPUExecutionProvider",`
`721`	`721`	`reduce_range=self.reduce_range,`
`722`	`722`	`**kwargs,`
`723`	`723`	`)`