Quantization of pytorch models

DineshRajan_T: <p><a class="mention" href="/u/ptrblck">@ptrblck</a><br>
How do I quantize my model to FP16 after training it normally in pytorch.</p>

Zafar: <p>It depends on what kind of quantization you are talking about. The answer will be different if you are using PTQ or QAT.</p>
<h2>PTQ</h2>
<p>If you are just asking about post-training quantization (PTQ), you can simply cast the data types to <code>torch.float16</code>, like so:</p>
<pre><code class="lang-auto">from torch import nn

# This is for reference
model = nn.Linear(3, 3)
x = torch.randn(128, 3)
y = model(x)

# This is the quantized model
model_16 = copy.deepcopy(model)
model_16.weight = nn.Parameter(model_16.weight.half())
model_16.bias = nn.Parameter(model_16.bias.half())

# This is the quantized computation
x_16 = x.half()
y_16 = model_16(x_16)

# Check the quantization error:
x_norm = torch.norm(x - x_16)
y_norm = torch.norm(y - y_16)

print(f'x error norm: {x_norm:.2f}, x abs max: {x.abs().max():.2f}')
print(f'y error norm: {y_norm:.2f}, y abs max: {y.abs().max():.2f}')

x_sqnr = 20 * torch.log10(torch.norm(x) / (torch.norm(x - x_16)))
y_sqnr = 20 * torch.log10(torch.norm(y) / (torch.norm(y - y_16)))

print(f'x sqnr: {x_sqnr:.2f} dB')
print(f'y sqnr: {y_sqnr:.2f} dB')

## Results:
# x error norm: 0.00, x abs max: 3.13
# y error norm: 0.01, y abs max: 2.30
# x sqnr: 73.35 dB
# y sqnr: 67.21 dB
</code></pre>
<h2>QAT</h2>
<p>For QAT, please, refer to the <a href="https://pytorch.org/docs/stable/quantization.html" class="inline-onebox" rel="noopener nofollow ugc">Quantization — PyTorch 1.8.0 documentation</a> (search for the “fakequantize” and “fake quantization”.</p>

jerryzh168: <p>we are providing a way to do fp16 static quantization in fx graph mode quantization as well, it is ready in master. you can find an example of related tests:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/master/test/quantization/test_quantize_fx.py#L2145" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/test/quantization/test_quantize_fx.py#L2145" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/master/test/quantization/test_quantize_fx.py#L2145</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="2135" style="counter-reset: li-counter 2134 ;">
<li>            # weight</li>
<li>            ns.call_method("to"): 1 if is_reference else 0</li>
<li>        }</li>
<li>        self.checkGraphModeFxOp(</li>
<li>            model, data, QuantType.DYNAMIC, qlinear_fun,</li>
<li>            is_reference=is_reference,</li>
<li>            custom_qconfig_dict={"": float16_dynamic_qconfig},</li>
<li>            prepare_expected_node_occurrence=prepare_node_occurrence,</li>
<li>            expected_node_occurrence=convert_node_occurrence)</li>
<li>
<li class="selected">def test_linear_static_fp16(self):</li>
<li>    class FuncLinear(torch.nn.Module):</li>
<li>        def __init__(self, use_bias, has_relu, f_relu):</li>
<li>            super(FuncLinear, self).__init__()</li>
<li>            self.w = torch.randn(4, 30)</li>
<li>            self.b = torch.randn(4)</li>
<li>            self.use_bias = use_bias</li>
<li>            if has_relu:</li>
<li>                if f_relu:</li>
<li>                    self.relu = F.relu</li>
<li>                else:</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>