Fixed pruning and distillation bug and update examples

Signed-off-by: Cheng, Penghui <[email protected]>

Author: PenghuiCheng

docs/source/dataloader.md

@@ -100,7 +100,7 @@ calib_data = mx.io.ImageRecordIter(path_imgrec=dataset,
                                    ctx=args.ctx,
                                    **combine_mean_std)
 
-from neural_compressor import Quantization, common
+from neural_compressor.experimental import Quantization, common
 quantizer = Quantization('conf.yaml')
 quantizer.model = fp32_model
 quantizer.calib_dataloader = calib_data

docs/source/dataset.md

@@ -96,7 +96,7 @@ class Dataset(object):
 After defining the dataset class, pass it to the quantizer:
 
 ```python
-from neural_compressor import Quantization, common
+from neural_compressor.experimental import Quantization, common
 quantizer = Quantization(yaml_file)
 quantizer.calib_dataloader = common.DataLoader(dataset) # user can pass more optional args to dataloader such as batch_size and collate_fn
 quantizer.model = graph

examples/.config/model_params_pytorch.json

@@ -220,91 +220,61 @@
       "model_src_dir": "nlp/huggingface_models/text-classification/quantization/ptq_static/fx",
       "dataset_location": "",
       "input_model": "/tf_dataset/pytorch/glue_data/base_weights/bert_MRPC_output",
-      "yaml": "conf.yaml",
-      "strategy": "basic",
-      "batch_size": 64,
-      "new_benchmark": false
+      "batch_size": 64
     },
     "bert_base_CoLA": {
       "model_src_dir": "nlp/huggingface_models/text-classification/quantization/ptq_static/fx",
       "dataset_location": "",
       "input_model": "/tf_dataset/pytorch/glue_data/base_weights/bert_CoLA_output",
-      "yaml": "conf.yaml",
-      "strategy": "basic",
-      "batch_size": 64,
-      "new_benchmark": false
+      "batch_size": 64
     },
     "bert_base_STS-B": {
       "model_src_dir": "nlp/huggingface_models/text-classification/quantization/ptq_static/fx",
       "dataset_location": "",
       "input_model": "/tf_dataset/pytorch/glue_data/base_weights/bert_STS-B_output",
-      "yaml": "conf.yaml",
-      "strategy": "basic",
-      "batch_size": 64,
-      "new_benchmark": false
+      "batch_size": 64
     },
     "bert_base_SST-2": {
       "model_src_dir": "nlp/huggingface_models/text-classification/quantization/ptq_static/fx",
       "dataset_location": "",
       "input_model": "/tf_dataset/pytorch/glue_data/base_weights/bert_SST-2_output",
-      "yaml": "conf.yaml",
-      "strategy": "basic",
-      "batch_size": 64,
-      "new_benchmark": false
+      "batch_size": 64
     },
     "bert_base_RTE": {
       "model_src_dir": "nlp/huggingface_models/text-classification/quantization/ptq_static/fx",
       "dataset_location": "",
       "input_model": "/tf_dataset/pytorch/glue_data/base_weights/bert_RTE_output",
-      "yaml": "conf.yaml",
-      "strategy": "basic",
-      "batch_size": 64,
-      "new_benchmark": false
+      "batch_size": 64
     },
     "bert_large_MRPC": {
       "model_src_dir": "nlp/huggingface_models/text-classification/quantization/ptq_static/fx",
       "dataset_location": "",
       "input_model": "/tf_dataset/pytorch/glue_data/weights/bert_MRPC_output",
-      "yaml": "conf.yaml",
-      "strategy": "basic",
-      "batch_size": 64,
-      "new_benchmark": false
+      "batch_size": 64
     },
     "bert_large_SQuAD": {
       "model_src_dir": "nlp/huggingface_models/question-answering/quantization/ptq_static/fx",
       "dataset_location": "",
       "input_model": "",
-      "yaml": "conf.yaml",
-      "strategy": "basic",
-      "batch_size": 64,
-      "new_benchmark": false
+      "batch_size": 64
     },
     "bert_large_QNLI": {
       "model_src_dir": "nlp/huggingface_models/text-classification/quantization/ptq_static/fx",
       "dataset_location": "",
       "input_model": "/tf_dataset/pytorch/glue_data/weights/bert_QNLI_output",
-      "yaml": "conf.yaml",
-      "strategy": "basic",
-      "batch_size": 64,
-      "new_benchmark": false
+      "batch_size": 64
     },
     "bert_large_RTE": {
       "model_src_dir": "nlp/huggingface_models/text-classification/quantization/ptq_dynamic/fx",
       "dataset_location": "",
       "input_model": "/tf_dataset/pytorch/glue_data/weights/bert_large_rte", 
-      "yaml": "conf.yaml",
-      "strategy": "basic",
-      "batch_size": 64,
-      "new_benchmark": false
+      "batch_size": 64
     },
     "bert_large_CoLA": {
       "model_src_dir": "nlp/huggingface_models/text-classification/quantization/ptq_static/fx",
       "dataset_location": "",
       "input_model": "/tf_dataset/pytorch/glue_data/weights/bert_CoLA_output",
-      "yaml": "conf.yaml",
-      "strategy": "basic",
-      "batch_size": 64,
-      "new_benchmark": false
+      "batch_size": 64
     },
     "dlrm": {
       "model_src_dir": "recommendation/dlrm/quantization/ptq/eager",
@@ -436,10 +406,7 @@
       "model_src_dir": "nlp/huggingface_models/text-classification/quantization/ptq_static/fx",
       "dataset_location": "/tf_dataset/pytorch/glue_data_new",
       "input_model": "/tf_dataset/pytorch/huggingface/language_translation_pt/distilbert_mrpc",
-      "yaml": "conf.yaml",
-      "strategy": "basic",
-      "batch_size": 64,
-      "new_benchmark": false
+      "batch_size": 64
     },
     "albert_base_MRPC": {
       "model_src_dir": "nlp/huggingface_models/text-classification/quantization/ptq_dynamic/eager",
@@ -463,10 +430,7 @@
       "model_src_dir": "nlp/huggingface_models/text-classification/quantization/ptq_static/fx",
       "dataset_location": "/tf_dataset/pytorch/glue_data_new",
       "input_model": "/tf_dataset/pytorch/huggingface/language_translation_pt/funnel_mrpc",
-      "yaml": "conf.yaml",
-      "strategy": "basic",
-      "batch_size": 64,
-      "new_benchmark": false
+      "batch_size": 64
     },
     "bart_WNLI(rm)": {
       "model_src_dir": "eager/huggingface_models",
@@ -697,10 +661,7 @@
       "model_src_dir": "nlp/huggingface_models/text-classification/quantization/qat/fx",
       "dataset_location": "/tf_dataset/pytorch/glue_data/MRPC/",
       "input_model": "/tf_dataset2/models/pytorch/bert_model",
-      "yaml": "conf_qat.yaml",
-      "strategy": "basic",
-      "batch_size": 8,
-      "new_benchmark": false
+      "batch_size": 8
     },
     "wide_resnet101_2_fx": {
       "model_src_dir": "oob_models/gen-efficientnet-pytorch",

examples/helloworld/tf_example4/test.py

@@ -2,7 +2,7 @@
 import time
 import numpy as np
 from tensorflow import keras
-from neural_compressor.experimental import Quantization,  common
+from neural_compressor.experimental import Quantization, common
 
 def main():
 

examples/helloworld/tf_example6/README.md

@@ -64,7 +64,7 @@ python test.py --benchmark
 ### 6. Introduction
 * We only need to add the following lines for quantization to create an int8 model.
 ```python
-    from neural_compressor import Quantization
+    from neural_compressor.experimental import Quantization
     quantizer = Quantization('./conf.yaml')
     quantized_model = quantizer('./mobilenet_v1_1.0_224_frozen.pb')
     tf.io.write_graph(graph_or_graph_def=quantized_model,

examples/helloworld/tf_example6/test.py

@@ -9,7 +9,7 @@ def main():
     args = arg_parser.parse_args()
 
     if args.tune:
-        from neural_compressor import Quantization
+        from neural_compressor.experimental import Quantization
         quantizer = Quantization('./conf.yaml')
         quantized_model = quantizer("./mobilenet_v1_1.0_224_frozen.pb")
         tf.io.write_graph(graph_or_graph_def=quantized_model,

examples/pytorch/nlp/huggingface_models/text-classification/quantization/ptq_dynamic/fx/conf.yaml

@@ -528,6 +528,7 @@ def benchmark(model):
         conf = QuantizationAwareTrainingConfig()
         compression_manager = prepare_compression(model, conf)
         compression_manager.callbacks.on_train_begin()
+        trainer.model = compression_manager.model
         trainer.train()
         compression_manager.callbacks.on_train_end()