From d6149aac01c8142f6f5ffc18c03433c82f44150c Mon Sep 17 00:00:00 2001
From: Yi Liu <yi4.liu@intel.com>
Date: Wed, 16 Oct 2024 14:02:29 +0800
Subject: [PATCH] Update the PT2E CV example (#2032)

Signed-off-by: yiliu30 <yi4.liu@intel.com>
---
 examples/3.x_api/pytorch/cv/static_quant/README.md       | 3 ++-
 examples/3.x_api/pytorch/cv/static_quant/main.py         | 9 ++++++---
 .../3.x_api/pytorch/cv/static_quant/run_benchmark.sh     | 4 ++--
 examples/3.x_api/pytorch/cv/static_quant/run_quant.sh    | 6 +++++-
 4 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/examples/3.x_api/pytorch/cv/static_quant/README.md b/examples/3.x_api/pytorch/cv/static_quant/README.md
index 172f8b0e12f..8a25bc251da 100644
--- a/examples/3.x_api/pytorch/cv/static_quant/README.md
+++ b/examples/3.x_api/pytorch/cv/static_quant/README.md
@@ -14,7 +14,8 @@ This implements quantization of popular model architectures, such as ResNet on t
 To quant a model and validate accaracy, run `main.py` with the desired model architecture and the path to the ImageNet dataset:
 
 ```bash
-python main.py -a resnet18 [imagenet-folder with train and val folders] -q -e
+export ImageNetDataPath=/path/to/imagenet
+python main.py $ImageNetDataPath --pretrained -a resnet18 --tune --calib_iters 5
 ```
 
 
diff --git a/examples/3.x_api/pytorch/cv/static_quant/main.py b/examples/3.x_api/pytorch/cv/static_quant/main.py
index 3d7af7827e3..655723221f0 100644
--- a/examples/3.x_api/pytorch/cv/static_quant/main.py
+++ b/examples/3.x_api/pytorch/cv/static_quant/main.py
@@ -46,6 +46,7 @@
 parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
                     metavar='W', help='weight decay (default: 1e-4)',
                     dest='weight_decay')
+parser.add_argument('--w_granularity', default="per_channel", type=str, choices=["per_channel", "per_tensor"], help='weight granularity')
 parser.add_argument('-p', '--print-freq', default=10, type=int,
                     metavar='N', help='print frequency (default: 10)')
 parser.add_argument('--resume', default='', type=str, metavar='PATH',
@@ -179,7 +180,7 @@ def eval_func(model):
     
     if args.tune:
         from neural_compressor.torch.export import export
-        from neural_compressor.torch.quantization import prepare, convert, get_default_static_config
+        from neural_compressor.torch.quantization import prepare, convert, StaticQuantConfig
 
         # Prepare the float model and example inputs for exporting model
         x = torch.randn(args.batch_size, 3, 224, 224).contiguous(memory_format=torch.channels_last)
@@ -188,7 +189,7 @@ def eval_func(model):
         # Specify that the first dimension of each input is that batch size
         from torch.export import Dim
         print(args.batch_size)
-        batch = Dim("batch", min=16)
+        batch = Dim("batch")
     
         # Specify that the first dimension of each input is that batch size
         dynamic_shapes = {"x": {0: batch}}
@@ -196,7 +197,7 @@ def eval_func(model):
         # Export eager model into FX graph model
         exported_model = export(model=model, example_inputs=example_inputs, dynamic_shapes=dynamic_shapes)
         # Quantize the model
-        quant_config = get_default_static_config()
+        quant_config = StaticQuantConfig(w_granularity=args.w_granularity)
         
         prepared_model = prepare(exported_model, quant_config=quant_config)
         # Calibrate
@@ -233,7 +234,9 @@ def eval_func(model):
             new_model = opt_model
         else:
             new_model = model
+            # For fair comparison, we also compile the float model
             new_model.eval()
+            new_model = torch.compile(new_model)
         if args.performance:
             benchmark(val_loader, new_model, args)
             return
diff --git a/examples/3.x_api/pytorch/cv/static_quant/run_benchmark.sh b/examples/3.x_api/pytorch/cv/static_quant/run_benchmark.sh
index 6f6b69c35df..cff26385b50 100644
--- a/examples/3.x_api/pytorch/cv/static_quant/run_benchmark.sh
+++ b/examples/3.x_api/pytorch/cv/static_quant/run_benchmark.sh
@@ -79,7 +79,7 @@ function run_benchmark {
         python main.py \
                 --pretrained \
                 -a resnet18 \
-                -b 30 \
+                -b ${batch_size} \
                 --tuned_checkpoint ${tuned_checkpoint} \
                 ${dataset_location} \
                 ${extra_cmd} \
@@ -89,7 +89,7 @@ function run_benchmark {
                 main.py \
                 --pretrained \
                 -a resnet18 \
-                -b 30 \
+                -b ${batch_size} \
                 --tuned_checkpoint ${tuned_checkpoint} \
                 ${dataset_location} \
                 ${extra_cmd} \
diff --git a/examples/3.x_api/pytorch/cv/static_quant/run_quant.sh b/examples/3.x_api/pytorch/cv/static_quant/run_quant.sh
index 1f4588e933c..186221cfd8a 100644
--- a/examples/3.x_api/pytorch/cv/static_quant/run_quant.sh
+++ b/examples/3.x_api/pytorch/cv/static_quant/run_quant.sh
@@ -10,6 +10,7 @@ function main {
 
 # init params
 function init_params {
+  batch_size=16
   tuned_checkpoint="saved_results"
   for var in "$@"
   do
@@ -22,6 +23,9 @@ function init_params {
       ;;
       --input_model=*)
           input_model=$(echo $var |cut -f2 -d=)
+      ;;
+      --batch_size=*)
+          batch_size=$(echo $var |cut -f2 -d=)
       ;;
        --output_model=*)
            tuned_checkpoint=$(echo $var |cut -f2 -d=)
@@ -44,7 +48,7 @@ function run_tuning {
             --pretrained \
             -t \
             -a resnet18 \
-            -b 30 \
+            -b ${batch_size} \
             --tuned_checkpoint ${tuned_checkpoint} \
             ${dataset_location}
 }