Selective Quantization

thyeros: <p>Hi, I like to selectively quantize layers as some layers in my project just serve as a regularizer. So, I tried a few ways and got confused with the following results.</p>
<pre><code class="lang-auto">class LeNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Linear(28 * 28, 10)    
        self.relu1 = nn.ReLU(inplace=True)

    def forward(self, x):             
        return self.relu1(self.l1(x.view(x.size(0), -1)))  
</code></pre>
<p><strong>1. selective qconfig assignment and top level transform</strong></p>
<pre><code class="lang-auto">model = LeNet()
model.l1.qconfig = torch.quantization.get_default_qat_qconfig()
torch.quantization.prepare_qat(model, inplace=True)
print(model)
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/f/1f4ff16006b06edf63ef6d7932aa821a13cb9ecc.png" data-download-href="https://discuss.pytorch.org/uploads/default/1f4ff16006b06edf63ef6d7932aa821a13cb9ecc" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/1/f/1f4ff16006b06edf63ef6d7932aa821a13cb9ecc.png" alt="image" data-base62-sha1="4t04Rm0c1W5XOq7x43Z1BZaliTq" width="690" height="214" data-dominant-color="EEEEEE"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1141×355 29.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<strong>2. selective qconfig assignment and selective transform</strong></p>
<pre><code class="lang-auto">model2 = LeNet()
model2.l1.qconfig = torch.quantization.get_default_qat_qconfig()
torch.quantization.prepare_qat(model2.l1, inplace=True)
print(model2)
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/a/2a79778cf025eb3a23c4cdf102ee7df1aabd9a71.png" data-download-href="https://discuss.pytorch.org/uploads/default/2a79778cf025eb3a23c4cdf102ee7df1aabd9a71" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/2/a/2a79778cf025eb3a23c4cdf102ee7df1aabd9a71.png" alt="image" data-base62-sha1="63KikgiZ80BnkYK9XTJLP2nJPUd" width="690" height="148" data-dominant-color="F1F1F1"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1139×245 16.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>You can see that the 2nd case doesn’t have  <strong>(weight_fake_quant): FakeQuantize</strong>. Is this a correct behavior? Shouldn’t both yield the same transformed model?<br>
Also, if there is a better way to do selective quantization (like different bits, quant vs no quant), please advise.</p>

Vasiliy_Kuznetsov: <p>You can use the <code>model.layer.qconfig = None</code> syntax to turn off quantization for a layer and all of its children.  Please feel free to see <a href="https://pytorch.org/docs/stable/quantization.html#model-preparation-for-quantization" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/quantization.html#model-preparation-for-quantization</a> for more context.</p>

thyeros: <p>Right, so in my both examples, qconifg is set for only ‘l1’ layer. But, the question is ‘should I always do prepare_qat at the top level?’. why does calling prepare_qat for a particular layer one at a time yields difference? Any explanation/insight would be appreciated</p>

hx89: <p>I think the reason is prepare_qat() calls convert(), which doesn’t convert the root module, so if you print the type of l1, in case 2 model2.l1 is root module thus not converted and still has type &lt;class ‘torch.nn.modules.linear.Linear’&gt; which doesn’t have weight_fake_quant attribute, while model.l1 is type &lt;class ‘torch.nn.qat.modules.linear.Linear’&gt; which has weight_fake_quant attribute.</p>

thyeros: <p>Sounds you’re right: it doesn’t convert the very model given to prepare_qat()</p>
<pre><code class="lang-auto">print(type(model2.l1))
&lt;class 'torch.nn.modules.linear.Linear'&gt;
</code></pre>