Fusing a QAT model post-training

username: GangnamStyle
cooked: <p>Hello, I have a QAT model that I have already trained. After training though, I would like to fuse the conv, bn, and relu layers of the model. However, after doing this, I seem to lose the quantization scale for the resulting ConvReLU2d layer. I’ve created a minimial example.</p>
<pre data-code-wrap="python"><code class="lang-python">class Model(th.nn.Module):
    def __init__(self):
        super().__init__()
        self.quant = th.quantization.QuantStub()
        self.conv = th.nn.Conv2d(1, 3, 3, bias=False)
        self.bn = th.nn.BatchNorm2d(3)
        self.relu = th.nn.ReLU()
        self.dequant = th.quantization.DeQuantStub()

    def forward(self, x):
        x = self.quant(x)
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        x = self.dequant(x)
        return x

# Initialize model
m = Model()

# Set qconfig
m.qconfig = th.ao.quantization.get_default_qconfig("fbgemm")

# Train
x = th.rand(64, 1, 32, 30)
pm = th.quantization.prepare_qat(m.train())
y = pm(x)
</code></pre>
<p>At this point, I can retrieve the quantized model as normal:</p>
<pre data-code-wrap="python"><code class="lang-python">qm = th.ao.quantization.convert(pm.eval())
</code></pre>
<p>And I can access the quantized weights of the convolutional layer as normal:</p>
<pre data-code-wrap="python"><code class="lang-python">qm.conv.weight().int_repr()
</code></pre>
<p>However, what I really want to do is fuse the conv, bn and relu layers together. I’ve tried the following:</p>
<pre data-code-wrap="python"><code class="lang-python">fused_pm = th.ao.quantization.fuse_modules(pm, [['conv','bn','relu']])
</code></pre>
<p>This results in the following error:</p>
<pre data-code-wrap="python"><code class="lang-python">AssertionError: did not find fuser method for: (&lt;class 'torch.ao.nn.qat.modules.conv.Conv2d'&gt;, &lt;class 'torch.nn.modules.batchnorm.BatchNorm2d'&gt;, &lt;class 'torch.nn.modules.activation.ReLU'&gt;) 
</code></pre>
<p>I receive a similar error if I instead try to fuse the quantized model (qm).</p>
<p>Is there a way to fuse a QAT model after training and access the updated quantized version? Thank you.</p>

username: jerryzh168
cooked: <aside class="quote no-group quote-modified" data-username="GangnamStyle" data-post="1" data-topic="205638">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/gangnamstyle/48/75622_2.png" class="avatar"> GangnamStyle:</div>
<blockquote>
<p>when I access the weights of the resulting conv (whose weights should have been updated based on the batchnorm), it’s just a normal tensor, not a quantized tensor, so I can’t call <code>int_repr()</code> on it.</p>
</blockquote>
</aside>
<p>the order for QAT with fused module should be:</p>
<ol>
<li>fuse_modules</li>
<li>prepare_qat</li>
<li>convert</li>
</ol>
<p>can you try this?</p>

username: GangnamStyle
cooked: <p>In my initial example, I realized I was using <code>prepare</code> instead of <code>prepare_qat</code>. I’ve updated my question since. At this point I have my prepared and trained model.  I tried fusing by doing this:</p>
<pre data-code-wrap="python"><code class="lang-python">fused_pm = th.quantization.fuse_modules(pm.eval(), [["conv", "bn", "relu"]])
</code></pre>
<p>However, I get the following error:</p>
<pre data-code-wrap="python"><code class="lang-python">AssertionError: did not find fuser method for: (&lt;class 'torch.ao.nn.qat.modules.conv.Conv2d'&gt;, &lt;class 'torch.nn.modules.batchnorm.BatchNorm2d'&gt;, &lt;class 'torch.nn.modules.activation.ReLU'&gt;) 
</code></pre>
<p>Is fusion not supported for QAT versions of conv and bn? Or is there a different function I should be using to fuse?</p>

username: jerryzh168
cooked: <p>you should apply fusion before prepare_qat</p>

username: jerryzh168
cooked: <p>or you could try out our new API: <a href="https://pytorch.org/tutorials/prototype/pt2e_quant_qat.html#prototype-pytorch-2-export-quantization-aware-training-qat" class="inline-onebox" rel="noopener nofollow ugc">(prototype) PyTorch 2 Export Quantization-Aware Training (QAT) — PyTorch Tutorials 2.3.0+cu121 documentation</a> all these details will be taken care of for you</p>