FX mode static_quantization for YOLOv7

RonaldYuren: <p>I use the FX_mode post-training-static quantization on YOLOv7，<br>
but I always see the error about the control-flow (which is happend at yolov7/models/yolo.py  , that function is fuseforward()， and I don’t know how to solve it)</p>
<p>you can see the code below:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/WongKinYiu/yolov7/blob/main/models/yolo.py">
  <header class="source">

      <a href="https://github.com/WongKinYiu/yolov7/blob/main/models/yolo.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/WongKinYiu/yolov7/blob/main/models/yolo.py" target="_blank" rel="noopener nofollow ugc">WongKinYiu/yolov7/blob/main/models/yolo.py</a></h4>


      <pre><code class="lang-py">import argparse
import logging
import sys
from copy import deepcopy

sys.path.append('./')  # to run '$ python *.py' files in subdirectories
logger = logging.getLogger(__name__)
import torch
from models.common import *
from models.experimental import *
from utils.autoanchor import check_anchor_order
from utils.general import make_divisible, check_file, set_logging
from utils.torch_utils import time_synchronized, fuse_conv_and_bn, model_info, scale_img, initialize_weights, \
    select_device, copy_attr
from utils.loss import SigmoidBin

try:
    import thop  # for FLOPS computation
except ImportError:
    thop = None
</code></pre>



  This file has been truncated. <a href="https://github.com/WongKinYiu/yolov7/blob/main/models/yolo.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>My pytorch version is 2.1.0, and I use the code in pytorch 2.1.0 document. to do quantization<br>
Please help me to solve it, Thank you!</p>

HDCharles: <p>that’s the main issue with fx quantization, you’re model needs to be symbolically traceable. You’ll need to get rid of the control flow or try some other quantization method.</p>
<p>This error is documented in the quantization docs here: <a href="https://pytorch.org/docs/stable/quantization.html#symbolic-trace-error-when-using-fx-graph-mode-quantization" class="inline-onebox" rel="noopener nofollow ugc">Quantization — PyTorch 2.2 documentation</a> with links to additional resources for making models symbolically traceable.</p>

RonaldYuren: <p>I have corrected this error by modified the model’s code.<br>
But I have encountered the “precision loss” after I using the Post-traing FX-mode  static quantization during my inference job.<br>
May I ask you “how to find the problem” or “where  to modified” that I can continue my quantization job?</p>

jerryzh168: <aside class="quote no-group" data-username="RonaldYuren" data-post="3" data-topic="198616">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ronaldyuren/48/66030_2.png" class="avatar"> RonaldYuren:</div>
<blockquote>
<p>or “where to modified” that I can continue my quantization job?</p>
</blockquote>
</aside>
<p>here is a guide for numerical debugging: <a href="https://pytorch.org/docs/stable/quantization-accuracy-debugging.html" class="inline-onebox" rel="noopener nofollow ugc">Quantization Accuracy Debugging — PyTorch 2.2 documentation</a>, but we may have updated tools later</p>

RonaldYuren: <p>Does it work when I am using pytorch 2.1.0?</p>

jerryzh168: <p>yeah it should work on 2.1</p>

RonaldYuren: <p>Thx a lot!!<br>
I am not sure about the fourth advice that is " consider QAT when trying PTQ",  that means I need to use QAT(on my training step ) and then implement PTQ?  or after I implement PTQ, then I use QAT just like finetuning?</p>

RonaldYuren: <p>Hi, I am trying using my “static_quantize_weights.pt” to do QAT(or you can call finetuning)；<br>
But it come across the problem that quantized-weight doesn’t have the same key compare with the original model.</p>
<p>the ERROR message is like below:</p>
<p>Traceback (most recent call last):<br>
File “Static_quantize5.py”, line 200, in <br>
train(hyp, opt = opt,device = torch.device(‘cpu’), tb_writer = tb_writer)<br>
File “/home/gene/YuRen/yolov7/train.py”, line 102, in train<br>
state_dict = ckpt[‘model’].float().state_dict()  # to FP32<br>
KeyError: ‘model’</p>
<p>How could I do next?</p>

jerryzh168: <p>it looks like it’s just <code>ckpt</code> does not have <code>model</code> as a key, so maybe you can checkout what does it contain and see what to do next</p>

RonaldYuren: <p><strong>Hi ,</strong><br>
<strong>I am using the code below to do quantization:</strong><br>
###############<br>
def saveModel(model):<br>
torch.save(model.state_dict() , ‘static_quantized_weights.pt’)</p>
<p>def static_quantize(weight_path, qconfig_mapping):<br>
weights, imgsz, cali_dataset = opt.weights, opt.img_size, opt.datasets</p>
<pre><code>#   Load Model
# 
device = torch.device("cpu")
# Load model  

ckpt = torch.load(weights[0])

model = attempt_load(weights, map_location=device)

stride = int(model.stride.max())  # model stride
imgsz = check_img_size(imgsz, s=stride)  # check img_size

dataset_processing = LoadImages(cali_dataset, img_size=imgsz, stride=stride)

model_to_quantize = copy.deepcopy(model)

example_input = torch.randn(1, 3, 1024, 1024)
qconfig_mapping = get_default_qconfig_mapping("fbgemm")  # better for intel CPU
model_to_quantize.eval()

# prepare
model_prepared = quantize_fx.prepare_fx(model_to_quantize, qconfig_mapping, example_input)
calibration(model_prepared, dataset_processing)
# quantize
model_quantized = quantize_fx.convert_fx(model_prepared)
saveModel(model_quantized)

loaded_model = torch.load('static_quantized_weights.pt', map_location='cpu')
print("Static Quantization Success!!!!")
</code></pre>
<p>##########################</p>
<p><strong>and I using the code below to load quantized-model:</strong><br>
###################<br>
ckpt = torch.load(weights, map_location=device)  # load checkpoint<br>
model_ckpt = attempt_load(weights, map_location=device) # you will get the weight from<br>
ckpt[‘model’]<br>
exclude = [‘anchor’] if (opt.cfg or hyp.get(‘anchors’)) and not opt.resume else <span class="chcklst-box fa fa-square-o fa-fw"></span>  # exclude keys</p>
<pre><code>        quantized_weights = torch.load(opt.quantization_weights, map_location='cpu')
        example_input = torch.randn(1, 3, 1024, 1024)
        qconfig_mapping = get_default_qconfig_mapping("fbgemm")
        model_ckpt.to(device)
        
        anchor_grid_modules2 = find_anchor_grid_modules(model_ckpt)
        grid_file_path2 = "Original_anchor_grid_modules.txt"
        write_anchor_grid_modules_to_file(anchor_grid_modules2, grid_file_path2)
        
        with open('parameter_groups_Originalmodel.txt', 'w') as f:
            for name, param in model_ckpt.named_parameters():   # OK  have value
                f.write(f'{name}\n')
        with open('module_names.txt', 'w') as f:
            # Iterate over named modules
            for name, module in model_ckpt.named_modules():
                # Write module name to file
                f.write(f'{module}\n')
        # prepare
        model_prepared = quantize_fx.prepare_fx(model_ckpt, qconfig_mapping, example_input)
        
        # calibrate (not shown)
        
        # quantize
        model = quantize_fx.convert_fx(model_prepared)
        
        model.load_state_dict(quantized_weights, strict=False)
        model.to(device) 
        with open('model_state_dict_keys.txt', 'w') as f:
            for key in model.state_dict().keys():
                f.write(f'{key}\n')
        with open('quantized_model_state_dict.txt', 'w') as f:   # OK have value
            count = 0
            for key, value in model.state_dict().items():
                f.write(f'{key}: {value}\n')
                count += 1
                if count &gt;= 30:
                    break
        with open('parameter_groups_Qmodel.txt', 'w') as f:    # no value
            for idx, param in enumerate(model.parameters()):
                f.write(f'Parameter {idx}: {param.shape}\n')
        with open('Quan_module_names.txt', 'w') as f:    # OK have value
            # Iterate over named modules
            for name, module in model.named_modules():
                # Write module name to file
                f.write(f'{name}\n')
</code></pre>
<p>###############<br>
<strong>Just like above,   the model_ckpt has value of model_ckpt…named_parameters()  &amp; .named_modules()  &amp; .state_dict()  ；</strong><br>
<strong>But the model (which is load the quantized-weights)  doesn’t have value for model.named_parameters()</strong></p>
<p>I am not sure what happend to my loading quantized-model<br>
Do you have any idea, please?</p>

HDCharles: <aside class="quote no-group" data-username="RonaldYuren" data-post="10" data-topic="198616">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ronaldyuren/48/66030_2.png" class="avatar"> RonaldYuren:</div>
<blockquote>
<pre><code class="lang-auto">        with open('model_state_dict_keys.txt', 'w') as f:
</code></pre>
</blockquote>
</aside>
<p>have you looked at <a href="https://pytorch.org/docs/stable/quantization.html#saving-and-loading-quantized-models" class="inline-onebox" rel="noopener nofollow ugc">Quantization — PyTorch 2.3 documentation</a>?</p>
<p>the other thing you could try is to do something along the lines:</p>
<pre><code class="lang-auto">model_prep = preapre_fx(model)
calibration(model_prep)
save_state_dict(model_prep.get_state_dict(), "prep_state_dict.pt")
....etc....
model = get_new_model()
model_prep = prepare_fx(model)
prep_state_dict = torch.load("prep_state_dict.pt")
model_prep.load_state_dict(prep_state_dict)
model_quantized = convert_fx(model_prep)
</code></pre>

RonaldYuren: <p>Oh… I have not read the 2.3 documentation because My computer is only suitable for pytorch2.1<br>
I will give a try today, Thx!!!</p>

RonaldYuren: <p>So… I think I should Load the prepare-model_stateDic.pt for my original model then go convert_fx(),<br>
And then… I also need to implemtent the code  “model.load_state_dict(quantized_weights, strict=False)”  to load my quantized-weights ?</p>

HDCharles: <p>you would do one or the other.</p>
<p>if you want to load the prepared state_dict, you would do the steps up to calibration, then after calibration you save the state_dict. to load, you would create a fresh model, apply the same prepare_fx setup, then load the prepared state_dict, then convert the model to get the final quantized model. In this case what you’d be saving/loading is the calibration statistics we collected so convert_fx can be applied correctly.</p>
<p>if you want to load the converted state dict, you would do the steps up to convert_fx and then save that state_dict. to load, you would create a fresh model, prepare and convert it, then load the state_dict. In this case you’d be saving/loading the actual quantized weights. However there are sometimes some weird serialization bugs.</p>
<p>either should work.</p>

RonaldYuren: <p>About the first way " save model’s state_dict when I finish the prepare &amp; Calibration , then implement the convert_fx when I am doing retraining ",<br>
but if I using the retraining model’s weights to do inference job , maybe some of the parameters(or model’s dictionary keys ) will disppear just like my problem before ?</p>
<p>I think I need to save the retrained quantized-model’s weights by saving state_dict(), or I have to use the torch.jit.save/load instead?<br>
I am not sure if it(torch.jit) will work for my pytorch 2.1</p>

RonaldYuren: <p>I think the " torch.jit " model type is not compatible to my code.</p>

Tangxiwei: <p>Can you tell me how to modified the model’s code for QAT ? I hope to hear from you soon.​ Thanks</p>