The packing format of quantized parameters after jitting

masahi: <p>Hi, following <a href="https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html" rel="nofollow noopener">the static quantization tutorial,</a>, I am trying to extract parameters of quantized, and jitted model. It seems after jitting, parameters are packed in a way that I don’t understand. For example, if I run the snippet below after the tutorial script, I get the output below.</p>
<pre><code class="lang-auto">input_size = (1, 3, 224, 224)
inp = np.random.randn(*input_size).astype("float32")
trace = torch.jit.trace(per_channel_quantized_model, torch.from_numpy(inp))
state_dict = trace.state_dict()
for (k, v) in state_dict.items():
    print(k, v.size())
</code></pre>
<pre><code class="lang-auto">features.0.0._packed_params torch.Size([128])
features.1.conv.0.0._packed_params torch.Size([128])
features.1.conv.1._packed_params torch.Size([128])
features.2.conv.0.0._packed_params torch.Size([128])
features.2.conv.1.0._packed_params torch.Size([128])
features.2.conv.2._packed_params torch.Size([128])
features.3.conv.0.0._packed_params torch.Size([128])
features.3.conv.1.0._packed_params torch.Size([128])
features.3.conv.2._packed_params torch.Size([128])
features.4.conv.0.0._packed_params torch.Size([128])
features.4.conv.1.0._packed_params torch.Size([128])
features.4.conv.2._packed_params torch.Size([128])
features.5.conv.0.0._packed_params torch.Size([128])
features.5.conv.1.0._packed_params torch.Size([128])
features.5.conv.2._packed_params torch.Size([128])
features.6.conv.0.0._packed_params torch.Size([128])
features.6.conv.1.0._packed_params torch.Size([128])
features.6.conv.2._packed_params torch.Size([128])
features.7.conv.0.0._packed_params torch.Size([128])
features.7.conv.1.0._packed_params torch.Size([128])
features.7.conv.2._packed_params torch.Size([128])
features.8.conv.0.0._packed_params torch.Size([128])
features.8.conv.1.0._packed_params torch.Size([128])
features.8.conv.2._packed_params torch.Size([128])
features.9.conv.0.0._packed_params torch.Size([128])
features.9.conv.1.0._packed_params torch.Size([128])
features.9.conv.2._packed_params torch.Size([128])
features.10.conv.0.0._packed_params torch.Size([128])
features.10.conv.1.0._packed_params torch.Size([128])
features.10.conv.2._packed_params torch.Size([128])
features.11.conv.0.0._packed_params torch.Size([128])
features.11.conv.1.0._packed_params torch.Size([128])
features.11.conv.2._packed_params torch.Size([128])
features.12.conv.0.0._packed_params torch.Size([128])
features.12.conv.1.0._packed_params torch.Size([128])
features.12.conv.2._packed_params torch.Size([128])
features.13.conv.0.0._packed_params torch.Size([128])
features.13.conv.1.0._packed_params torch.Size([128])
features.13.conv.2._packed_params torch.Size([128])
features.14.conv.0.0._packed_params torch.Size([128])
features.14.conv.1.0._packed_params torch.Size([128])
features.14.conv.2._packed_params torch.Size([128])
features.15.conv.0.0._packed_params torch.Size([128])
features.15.conv.1.0._packed_params torch.Size([128])
features.15.conv.2._packed_params torch.Size([128])
features.16.conv.0.0._packed_params torch.Size([128])
features.16.conv.1.0._packed_params torch.Size([128])
features.16.conv.2._packed_params torch.Size([128])
features.17.conv.0.0._packed_params torch.Size([128])
features.17.conv.1.0._packed_params torch.Size([128])
features.17.conv.2._packed_params torch.Size([128])
features.18.0._packed_params torch.Size([128])
quant.scale torch.Size([1])
quant.zero_point torch.Size([1])
classifier.1._packed_params._packed_params torch.Size([104])
</code></pre>
<p>I have no idea what is going on in this format and I have many questions. But for now let me ask you these:</p>
<ul>
<li>Is there a documentation of the packing format?</li>
<li>How can I extract the original floating point tensors along with scale and zero point? I confirmed that they are available before tracing.</li>
<li>Or even better, is there a way to prevent packing?</li>
<li>During tracing, where in the code base does this packing happen?</li>
</ul>
<p>I’m trying to translate jitted, quantized PyTorch model to TVM IR. For that I need floating point tensors with scale and zero point. That is the reason I’m asking here.</p>
<p>cc <a class="mention" href="/u/raghuramank100">@raghuramank100</a> <a class="mention" href="/u/jerryzh168">@jerryzh168</a></p>

masahi: <p>ok torch.ops.quantized.conv2d_unpack did the job.</p>

johnzhou1996: <p>Hello, I met the same problem. Could you show me the detail of the “torch.ops.quantized.conv2d_unpack”? And how to deal with classifier.1._packed_params?<br>
Thanks!</p>

masahi: <p>See the implementation in TVM I added:<br>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/apache/incubator-tvm/blob/06e9542ee0bfd014bd06a4dd4fdb3af9d2d29eb0/python/tvm/relay/frontend/qnn_torch.py#L50-L100" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/apache/incubator-tvm/blob/06e9542ee0bfd014bd06a4dd4fdb3af9d2d29eb0/python/tvm/relay/frontend/qnn_torch.py#L50-L100" target="_blank" rel="nofollow noopener">apache/incubator-tvm/blob/06e9542ee0bfd014bd06a4dd4fdb3af9d2d29eb0/python/tvm/relay/frontend/qnn_torch.py#L50-L100</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="50" style="counter-reset: li-counter 49 ;">
<li>def _unpack_quant_params(param_name, packed_params, unpack_func):</li>
<li>    # Torch stores quantized params in a custom packed format,</li>
<li>    # need to unpack and retrieve them as numpy arrays</li>
<li>    qweight, bias = unpack_func(packed_params)</li>
<li>    weight_np = qweight.dequantize().numpy()</li>
<li>
</li>
<li>    import torch</li>
<li>    if qweight.qscheme() == torch.per_tensor_affine:</li>
<li>        param = QNNParam(weight_np, bias, qweight.q_scale(),</li>
<li>                         int(qweight.q_zero_point()), param_name)</li>
<li>    else:</li>
<li>        scales = qweight.q_per_channel_scales().numpy()</li>
<li>        zero_points = qweight.q_per_channel_zero_points().numpy()</li>
<li>        # This is an assumption posed by QNN</li>
<li>        msg = "The values of zero points should be all zero for per channel"</li>
<li>        assert np.all(zero_points == 0), msg</li>
<li>        param = QNNParam(weight_np, bias, scales, 0, param_name)</li>
<li>
</li>
<li>    return param</li>
<li>
</li>
</ol></code></pre>

  This file has been truncated. <a href="https://github.com/apache/incubator-tvm/blob/06e9542ee0bfd014bd06a4dd4fdb3af9d2d29eb0/python/tvm/relay/frontend/qnn_torch.py#L50-L100" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>From the name <code>classifier.1._packed_params</code> I guess it comes from nn.Linear. In that case, you need to use <code>torch.ops.quantized.linear_unpack</code>.</p>

k.osama: <p>Hi I am working with a quantized model in C++, I wonder if I can parse the jitted model parameters like this in C++ ? I could not find any unpacking modules in  torch::jit::script::Module . I have trained and quantized my model in Python and loaded to C++. I am using version 1.6.0+</p>