Using quantizable model for normal training

username: Hackster
cooked: <p>I want to integrate quantization aware training (QAT) as an option into my training code. Is there anything that speaks against preparing models for QAT and training them normally, i.e. without QAT?</p>
<p>With preparation I mean adding QuantStub, FloatFunctional, and using already fused modules such as ConvReLU2d. As far as I dived into the source these modules only serve as placeholders for fake quantization modules when actually converting model for QAT and do not affect training.</p>
<p>Question comes up because I saw that torchvision includes two versions of models, a normal and a quantizable version. If not depending on backward compatibility this appears unnecessary for my own code. Is there something I oversee?</p>

username: andrewor
cooked: <p>Hi, by “quantizable” I suppose you mean the ones in torchvision specifically. These are needed for eager mode quantization, since the user needs to manually insert QuantStubs and DeQuantStubs like you mentioned. For QAT, these will be replaced by FakeQuantizes, which actually do change the numerics of training. That’s the point of QAT in the first place, which is to improve the accuracy of quantization by making the training process “aware” that the model will ultimately be quantized.</p>
<p>So my recommendation is the following. Either make “quantizable” versions of your model similar to torchvision, which uses eager mode quantization, or switch to FX graph mode quantization, where you don’t have to change a thing about your model and it’ll still be quantized automatically (with FakeQuantizes inserted for the QAT case). You can learn more about FX graph mode quantization here: <a href="https://pytorch.org/tutorials/prototype/fx_graph_mode_quant_guide.html" class="inline-onebox" rel="noopener nofollow ugc">(prototype) FX Graph Mode Quantization User Guide — PyTorch Tutorials 2.0.1+cu117 documentation</a>. Please feel free to let me know if there’s anything else I can clarify.</p>
<p>Best,<br>
-Andrew</p>

username: Hackster
cooked: <p>Hi Andrew,<br>
as my model structures are a bit complicated, I need to use eager quantization. What I meant with “making models quantizable” is including all necessary modules for QAT or PTQ already, whether using it or not.</p>
<p>For example, instead of defining my model as:</p>
<pre><code class="lang-auto">class Model(nn.Module):

    def __init__(self, nc_in: int=3):
        super().__init__()
        self.conv = nn.Conv2d(nc_in, 64, 3, 1, 1, bias=False)
        self.bnorm = nn.BatchNorm2d(64, affine=True)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        return x + self.relu(self.bnorm(self.conv(x)))
</code></pre>
<p>and then fusing it, adding QuantStub/DeQuantstub etc. later or even maintain an additional model only for QAT, I could simply define my model as</p>
<pre><code class="lang-auto">from torch.ao.nn.intrinsic import ConvBnReLU2d
from torch.ao.nn.quantized import FloatFunctional
class QuantizableModel(nn.Module):

    def __init__(self, nc_in: int=3):
        super().__init__()
        self.quant = QuantStub()
        self.conv_bnorm_relu = ConvBnReLU2d(
                nn.Conv2d(nc_in, 64, 3, 1, 1, bias=False),
                nn.BatchNorm2d(64, affine=True),
                nn.ReLU(inplace=True),
            )
        self.add = FloatFunctional()
        self.dequant = DeQuantStub()

    def forward(self, x):
        x = self.quant(x)
        x = self.add.add(x, self.conv_bnorm_relu (x))
        x = self.dequant(x)
        return x
</code></pre>
<p>and use it for training whether using QAT or not.<br>
Thus, I would not have to deal with fusing modules or adding QuantStubs as everything is already specifically defined for this type of model. All that was required would be adding a qconfig and calling prepare_qat</p>
<pre><code class="lang-auto">if use_qat:
    model.qconfig = torch.ao.quantization.get_default_qconfig('x86')
    torch.quantization.prepare_qat(model, inplace=True)
</code></pre>
<p>And in case of not using QAT nothing would change compared to the previous model.</p>
<p>Well, the latter is my actual question as I was uncertain whether there are any pitfalls, any other changes conducted in torch.ao.quantization.fuse_modules_qat, or in the fused placeholders that might be an obstacle</p>

username: jerryzh168
cooked: <aside class="quote no-group" data-username="Hackster" data-post="1" data-topic="183801">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/h/74df32/48.png" class="avatar"> Hackster:</div>
<blockquote>
<p>Question comes up because I saw that torchvision includes two versions of models, a normal and a quantizable version. If not depending on backward compatibility this appears unnecessary for my own code. Is there something I oversee?</p>
</blockquote>
</aside>
<p>I think Conv - bn fusion might introduce some slight numerical mismatches, but I don’t see it affecting the model accuracy in general.</p>

username: Hackster
cooked: <p>This is true if the modules are actually fused.<br>
But as far as I understand it, the <code>ConvBnReLU2d</code> module and all other modules from <code>torch.ao.nn.intrinsic</code> only act as placeholders for the actual fused modules. If you look at its source code and the source of its base class <code>_FusedModule</code>, you see that it’s just a Sequential module calling one after another.</p>
<p><a href="https://pytorch.org/docs/stable/_modules/torch/ao/nn/intrinsic/modules/fused.html#ConvBnReLU2d" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/_modules/torch/ao/nn/intrinsic/modules/fused.html#ConvBnReLU2d</a></p>
<p>The actual fusion of modules for QAT seems to happen in <code>torch.quantization.prepare_qat</code>, replacing the placeholders with the corresponding fused qat modules.</p>

username: jerryzh168
cooked: <p>oh yes, that’s correct I think. fuse_modules won’t change the numerics</p>