fix pt notebook benchmark (#1451)

Signed-off-by: Cheng, Zixuan <zixuan.cheng@intel.com>

Author: violetch24

examples/notebook/pytorch/Quick_Started_Notebook_of_INC_for_Pytorch.ipynb

@@ -285,7 +285,8 @@
     "from neural_compressor.config import PostTrainingQuantConfig, TuningCriterion\n",
     "tuning_criterion = TuningCriterion(max_trials=600)\n",
     "conf = PostTrainingQuantConfig(approach=\"static\", tuning_criterion=tuning_criterion)\n",
-    "q_model = fit(model, conf=conf, calib_dataloader=eval_dataloader, eval_func=eval_func)"
+    "q_model = fit(model, conf=conf, calib_dataloader=eval_dataloader, eval_func=eval_func)\n",
+    "q_model.save(\"./saved_results\")"
    ]
   },
   {
@@ -304,10 +305,10 @@
    "outputs": [],
    "source": [
     "# fp32 benchmark\n",
-    "!{sys.executable} benchmark.py --input_model ./pytorch_model.bin 2>&1|tee fp32_benchmark.log\n",
+    "!{sys.executable} benchmark.py 2>&1|tee fp32_benchmark.log\n",
     "\n",
     "# int8 benchmark\n",
-    "!{sys.executable} benchmark.py --input_model ./saved_results/best_model.pt 2>&1|tee int8_benchmark.log\n"
+    "!{sys.executable} benchmark.py --input_model saved_results 2>&1|tee int8_benchmark.log\n"
    ]
   }
  ],

examples/notebook/pytorch/benchmark.py

@@ -73,7 +73,7 @@ def compute_metrics(p: EvalPrediction):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(__doc__)
-    parser.add_argument("--input_model", type=str, required=True)
+    parser.add_argument("--input_model", type=str, required=False, default=None)
     args = parser.parse_args()
 
     # Initialize our Trainer
@@ -87,9 +87,23 @@ def compute_metrics(p: EvalPrediction):
     )
 
     eval_dataloader = trainer.get_eval_dataloader()
+    if eval_dataloader.batch_size is None:
+        def _build_inc_dataloader(dataloader):
+            class INCDataLoader:
+                __iter__ = dataloader.__iter__
+                def __init__(self) -> None:
+                    self.dataloader = dataloader
+                    self.batch_size = dataloader.total_batch_size
+            return INCDataLoader()
+        eval_dataloader = _build_inc_dataloader(eval_dataloader)
+    batch_size = eval_dataloader.batch_size
+
+    if args.input_model:
+        from neural_compressor.utils.pytorch import load
+        model = load(args.input_model, model, dataloader=eval_dataloader)
 
     b_conf = BenchmarkConfig(warmup=5,
                         iteration=100,
                         cores_per_instance=4,
                         num_of_instance=1)
-    benchmark.fit(args.input_model, b_conf, b_dataloader=eval_dataloader)
+    benchmark.fit(model, b_conf, b_dataloader=eval_dataloader)