Quantized model takes more time than fp32 model on CPU

KURI: <p>I used <code>fbgemm</code> as qconfig, and I checked that my cpu (Intel Xeon silver 4114) supports AVX2 operations. But my quantized model takes 3 times longer to inference than original fp32 model.</p>
<p>My original model is as follows,</p>
<pre><code class="lang-auto">#Original model
class NormUnet(nn.Module):
    def __init__(self, in_chans, out_chans, chans, num_pools):
        super().__init__()
        self.unet = Unet(
            in_chans=in_chans,
            out_chans=out_chans,
            chans=chans,
            num_pool_layers=num_pools
        )

    def norm(self, x):
        b, h, w = x.shape
        x = x.view(b, h * w)
        mean = x.mean(dim=1).view(b, 1, 1)
        std = x.std(dim=1).view(b, 1, 1)
        x = x.view(b, h, w)
        return (x - mean) / std, mean, std

    def unnorm(self, x, mean, std):
        return x * std + mean

    def forward(self, x):
        x, mean, std = self.norm(x)
        x = x.unsqueeze(1)
        x = self.unet(x)
        x = x.squeeze(1)
        x = self.unnorm(x, mean, std)
        return x
</code></pre>
<p>As <code>std</code> in <code>norm</code> method doesn’t support quantization, I only quantized unet, then put it back to original model.</p>
<pre><code class="lang-auto"># Load pre-trained model.
model = NormUnet(in_chans=1, out_chans=1, chans=128, num_pools=4)
model = torch.nn.parallel.DataParallel(model)
checkpoint = torch.load(model_path)
model.load_state_dict(checkpoint['model'])
model = model.module
model.to('cpu')

#Extract only unet part
unet_before_quantize = model.unet

#For quantization
class QuantizedUNet(nn.Module):
    def __init__(self, model_fp32):
        super(QuantizedUNet, self).__init__()
        self.quant = torch.quantization.QuantStub()
        self.dequant = torch.quantization.DeQuantStub()
        self.model_fp32 = model_fp32

    def forward(self, x):
        x = self.quant(x)
        x = self.model_fp32(x)
        x = self.dequant(x)
        return x

quantized_model = QuantizedUNet(model_fp32=unet_before_quantize)
quantization_config = torch.quantization.get_default_qconfig("fbgemm")
quantized_model.qconfig = quantization_config

torch.quantization.prepare(quantized_model, inplace=True)
</code></pre>
<p>Then I calibrate and quantize. After all these, Quantized model output is what I expected, but inference time gets about 3~4 times longer.</p>
<pre><code class="lang-auto">#I calibrated in between
quantized_model = torch.quantization.convert(quantized_model, inplace=True)

quantized_model.eval()

#I replace only unet block
model.unet = quantized_model
</code></pre>
<p>Where did I make mistakes?</p>
<p>I tried <code>torch.set_num_threads(1)</code>, and export to torchscript then load back. Couldn’t get faster than original fp32 model.</p>
<p>Helps will be really appreciated.</p>
<p>Thanks!</p>

Vasiliy_Kuznetsov: <p>Hi <a class="mention" href="/u/kuri">@KURI</a>, have you tried profiling both models to see which layers are the bottlenecks?  That would help narrow down the issue.  <a href="https://pytorch.org/tutorials/recipes/recipes/profiler.html" class="inline-onebox" rel="noopener nofollow ugc">PyTorch Profiler — PyTorch Tutorials 1.7.1 documentation</a> has an example.</p>

KURI: <p>Thanks for advice!</p>
<p>Unfortunately, profiler doesn’t work on quantized model. Its execution never ends. On fp32 model, profiler works well.</p>

Vasiliy_Kuznetsov: <p>Hi <a class="mention" href="/u/kuri">@KURI</a>, thanks for reporting, your code to call the quantization APIs looks good, and that sounds like a bug somewhere in PyTorch.  Is your model architecture shareable?  If it is, could you file a github issue with a repro and a description of your environment?  Then we could get someone on the fbgemm team to take a look.</p>