`torch.mean(keepdim=True)` with `quint8` tensor does not work with qnnpack backend

Alexander_Soare: <p>I’m new to quantization so I couldn’t figure out a way to easily reproduce this without going through the whole flow.</p>
<p>From <code>pdb</code> during a forward pass of a quantized model:</p>
<pre><code class="lang-auto">print(x.dtype)
# &gt;&gt; torch.quint8
print(x.shape)
# &gt;&gt; torch.Size([1, 40, 64, 384])
print(x.mean((2,3), keepdim=True).shape)
# &gt;&gt; torch.Size([1, 40])
</code></pre>
<p>This happens when I run the forward pass just after setting <code>torch.backends.quantized.engine = 'qnnpack'</code>.</p>
<p>If I do not set it, the forward pass runs fine, and 10x faster than the non-quantized version of my model (in other words, as expected)</p>
<p>Running this on Android causes the same issue.</p>

HDCharles: <p>I’m unable to reproduce:</p>
<pre><code class="lang-auto">import torch
import copy


# define a floating point model where some layers could benefit from QAT
class M(torch.nn.Module):
    def __init__(self):
        super(M, self).__init__()
        # QuantStub converts tensors from floating point to quantized
        self.quant = torch.quantization.QuantStub()
        self.conv = torch.nn.Conv2d(1, 1, 1)
        self.bn = torch.nn.BatchNorm2d(1)
        self.relu = torch.nn.ReLU()
        # DeQuantStub converts tensors from quantized to floating point
        self.dequant = torch.quantization.DeQuantStub()

    def forward(self, x):
        x = self.quant(x)
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        x = x.mean((2, 3), keepdim=True)
        x = self.dequant(x)
        return x


model_fp32 = M()
model_fp32.train()
model_fp32.qconfig = torch.quantization.get_default_qat_qconfig('qnnpack')
model_fp32_fused = torch.quantization.fuse_modules(model_fp32,
                                                   [['conv', 'bn', 'relu']])
model_fp32_prepared = torch.quantization.prepare_qat(model_fp32_fused)
x = torch.rand(2, 1, 224, 224)
model_fp32_prepared.eval()
model_fp32_prepared(x)
model_int8 = torch.quantization.convert(model_fp32_prepared)

y = model_int8(x)
print(y)
print(y.shape)
</code></pre>
<pre><code class="lang-auto">tensor([[[[0.]]],
        [[[0.]]]])
torch.Size([2, 1, 1, 1])
</code></pre>
<p>Do you have a better repro?</p>

tom: <p>Hello Alexander,</p>
<p>I can confirm this and took the liberty to file <a href="https://github.com/pytorch/pytorch/issues/58668" class="inline-onebox" rel="noopener nofollow ugc">QNNPACK mean with keepdim doesn't work · Issue #58668 · pytorch/pytorch · GitHub</a> . Thank you so much for reporting this with very precise repro information! This makes things much easier.</p>
<p>For reference: This illustrates the problem:</p>
<pre><code class="lang-python">torch.backends.quantized.engine = 'qnnpack'
print(torch.backends.quantized.engine, torch.quantize_per_tensor(torch.randn(5, 5, 5, 5), scale=0.2, zero_point=0, dtype=torch.quint8).mean((2,3), keepdim=True).shape)
torch.backends.quantized.engine = 'fbgemm'
print(torch.backends.quantized.engine, torch.quantize_per_tensor(torch.randn(5, 5, 5, 5), scale=0.2, zero_point=0, dtype=torch.quint8).mean((2,3), keepdim=True).shape)
</code></pre>
<p>Best regards</p>
<p>Thomas</p>

tom: <p>I’m happy to report that the issue linked above has been closed, so we should see nightlies that have the problem fixed. I don’t think it made it into 1.9, though, but I hope to make Raspberry Pi wheels with the fix soon enough.</p>
<p>Best regards</p>
<p>Thomas</p>