Use torchao for quantization (#65)

Author: cpuhrsch

.github/workflows/test_install.yml

@@ -28,4 +28,4 @@ jobs:
 
     - name: Test import
       run: |
-        python -c "from segment_anything_fast import dynamic_quant, sam_model_registry"
+        python -c "import segment_anything_fast; from segment_anything_fast import sam_model_registry"

experiments/eval_combo.py

@@ -337,7 +337,7 @@ def run(
         block.attn.use_rel_pos = use_rel_pos
 
     if compress == "dynamic_quant":
-        from segment_anything_fast.dynamic_quant import apply_dynamic_quant
+        from torchao.quantization import apply_dynamic_quant
         apply_dynamic_quant(predictor.model.image_encoder)
         inductorconfig.force_fuse_int_mm_with_mul = True
     elif compress == "static_quant":

experiments/p4d_results/results_bs8_vit_h.csv

@@ -1,9 +1,9 @@
 technique,time,sam_commit_name,pytorch_version,sam_model_type,batch_size,memory(MiB),memory(%),img_s(avg),batch_ms(avg)/batch_size,mIoU,use_compile,use_half,compress,epilogue_fusion_first,use_compile_decoder,use_nested_tensor,use_rel_pos,pad_input_image_batch,num_workers,num_batches,num_images,profile_path,memory_path
-fp32,33.77061091264089,default,2.2.0.dev20231024+cu121,vit_h,8,28806,71,2.7820335945039893,359.44928989194705,0.584173340367447,False,None,None,False,False,False,True,True,32,619,4952,None,None
-bf16,6.822473649183909,codesign,2.2.0.dev20231024+cu121,vit_h,8,14424,35,14.850424350893103,67.33814309756475,0.5809121174676433,False,torch.bfloat16,None,False,False,False,True,True,32,619,4952,None,None
-compile,7.968364950021108,codesign,2.2.0.dev20231024+cu121,vit_h,8,12358,30,19.69605657526638,50.77158446304247,0.5811320849834102,max-autotune,torch.bfloat16,None,False,False,False,True,True,32,619,4952,None,None
-SDPA,5.843019040425618,sdpa-decoder,2.2.0.dev20231024+cu121,vit_h,8,7947,19,21.92026495560376,45.61988653081299,0.581191777206921,max-autotune,torch.bfloat16,None,False,False,False,True,True,32,619,4952,None,None
-Triton,9.09047209819158,local-fork,2.2.0.dev20231024+cu121,vit_h,8,4550,11,22.874989934428537,43.71586623060877,0.5820036887609843,max-autotune,torch.bfloat16,None,False,False,False,True,True,32,619,4952,None,None
-NT,5.455243261655172,local-fork,2.2.0.dev20231024+cu121,vit_h,8,4550,11,23.206823845253847,43.09077393219044,0.5809004559961229,max-autotune,torch.bfloat16,None,False,False,True,True,True,32,619,4952,None,None
-int8,6.994769084453583,local-fork,2.2.0.dev20231026+cu121,vit_h,8,4167,10,24.87583443921619,40.19965651578395,0.5819033780783904,max-autotune,torch.bfloat16,dynamic_quant,False,False,True,True,True,32,619,4952,None,None
-sparse,5.597406772772471,local-fork,2.2.0.dev20231024+cu121,vit_h,8,7055,17,24.900183397177024,40.16034677533225,0.5289167514647479,max-autotune,torch.bfloat16,sparse,False,False,True,True,True,32,619,4952,None,None
+fp32,33.78124016523361,default,2.2.0.dev20231110+cu121,vit_h,8,28806,71,2.7815123804581523,359.51664534215973,0.584173340367447,False,None,None,False,False,False,True,True,32,619,4952,None,None
+bf16,6.781990921497345,codesign,2.2.0.dev20231110+cu121,vit_h,8,14424,35,14.949423825567571,66.89221013921143,0.5809121174676433,False,torch.bfloat16,None,False,False,False,True,True,32,619,4952,None,None
+compile,6.374551324049632,codesign,2.2.0.dev20231110+cu121,vit_h,8,12358,30,19.174280981156954,52.153194217959204,0.5809779984197878,max-autotune,torch.bfloat16,None,False,False,False,True,True,32,619,4952,None,None
+SDPA,5.68677978515625,sdpa-decoder,2.2.0.dev20231110+cu121,vit_h,8,7947,19,21.702906054763186,46.07677872616178,0.5810025119549971,max-autotune,torch.bfloat16,None,False,False,False,True,True,32,619,4952,None,None
+Triton,5.398879257837931,local-fork,2.2.0.dev20231110+cu121,vit_h,8,4550,11,23.150324132692813,43.19593947230325,0.5821156148026875,max-autotune,torch.bfloat16,None,False,False,False,True,True,32,619,4952,None,None
+NT,5.341391940911611,local-fork,2.2.0.dev20231110+cu121,vit_h,8,4550,11,23.571117205832582,42.42480283253413,0.580580762025661,max-autotune,torch.bfloat16,None,False,False,True,True,True,32,619,4952,None,None
+int8,5.64138038555781,local-fork,2.2.0.dev20231110+cu121,vit_h,8,4167,10,24.98278914178334,40.02755634387976,0.5821159807587736,max-autotune,torch.bfloat16,dynamic_quant,False,False,True,True,True,32,619,4952,None,None
+sparse,5.2578640898068745,local-fork,2.2.0.dev20231110+cu121,vit_h,8,7055,17,24.876148779272,40.19914854477989,0.5287772086382821,max-autotune,torch.bfloat16,sparse,False,False,True,True,True,32,619,4952,None,None

segment_anything_fast/dynamic_quant.py

@@ -13,6 +13,7 @@
         'pycocotools',
         'scipy',
         'scikit-image',
+        'torchao',
     ],
     include_package_data=True,
     package_data={