Update the PT2E CV example (#2032)

Signed-off-by: yiliu30 <[email protected]>

Author: yiliu30

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
 ```
 
 

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,15 +189,15 @@ 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}}
 
         # 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

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} \

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}
 }