RuntimeError: Unsupported qscheme: per_channel_affine during fx qat

username: Chame_call
cooked: <p>I use default qat qconfig for my model:<br>
<code>qconfig_mapping = get_default_qat_qconfig_mapping("x86")</code><br>
after I prepare model with the config<br>
<code>    model.backbone = prepare_qat_fx(model.backbone, qconfig_mapping, example_inputs)</code><br>
I see that for my first layer <code>per_channel_symmetric</code> is going to be applied.</p>
<pre><code class="lang-auto">&gt;&gt; dict(model.backbone.named_modules())['backbone.conv_stem'].qconfig.weight.p.keywords['qscheme']
&gt;&gt; torch.per_channel_symmetric
</code></pre>
<p>and then after training for one epoch I call <code>convert_fx(model.backbone)</code> and get the mentioned error.<br>
so there are actually two questions:</p>
<ol>
<li>I guess that during training the weight qscheme were changed for some reason but why?</li>
<li>why is it not able to apply the qscheme (<code>per_channel_affine</code>) especially considering that the scheme were changed automatically (I guess).<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/4/94e2c1724a8d4c7a6587df0f4bb11cb3ebc41c4e.png" data-download-href="https://discuss.pytorch.org/uploads/default/94e2c1724a8d4c7a6587df0f4bb11cb3ebc41c4e" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/4/94e2c1724a8d4c7a6587df0f4bb11cb3ebc41c4e_2_690x351.png" alt="image" data-base62-sha1="lf6udr3WNXDgyO8h6FSb2aDCOPI" width="690" height="351" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/4/94e2c1724a8d4c7a6587df0f4bb11cb3ebc41c4e_2_690x351.png, https://discuss.pytorch.org/uploads/default/optimized/3X/9/4/94e2c1724a8d4c7a6587df0f4bb11cb3ebc41c4e_2_1035x526.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/9/4/94e2c1724a8d4c7a6587df0f4bb11cb3ebc41c4e_2_1380x702.png 2x" data-dominant-color="401316"><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">1498×763 152 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></li>
</ol>

username: jerryzh168
cooked: <p>this is still <code>per_channel_symmetric</code>, it’s just in quantized tensors we don’t use it and convert it to <code>per_channel_affine</code> to simplify the kernel support (so you don’t have operation between per_channel_affine and per_channel_symmetric tensors etc.)</p>

username: Chame_call
cooked: <p>I see.<br>
But what about the main point? Why do we have the error for simple conv layer?</p>

username: jerryzh168
cooked: <p>sorry, not exactly sure, could you give me a minimal repro?</p>

username: Chame_call
cooked: <p>sure, here it is</p>
<pre><code class="lang-auto">from torch.ao.quantization import get_default_qat_qconfig_mapping
from torch.ao.quantization.quantize_fx import prepare_qat_fx, convert_fx
import timm
from tqdm import tqdm
import torch
from torch import nn

device = 'cuda' if torch.cuda.is_available() else 'cpu'

model = timm.create_model('mobilenetv3_large_100', pretrained=True, num_classes=3).to(device)

qconfig_mapping = get_default_qat_qconfig_mapping("x86")
example_inputs = torch.rand(size=(1,3,128,128)).to(device)

model = prepare_qat_fx(model, qconfig_mapping, example_inputs)

# train
def fake_train(model, n_samples=5, bs=64, device='cpu'):
    optimizer = torch.optim.AdamW(model.parameters(), lr=10e-4)
    loss = nn.CrossEntropyLoss()
    
    model.train()
    for _ in tqdm(range(n_samples)):
        x = torch.rand(size=(bs,3,128,128)).to(device)
        y = torch.empty(bs, dtype=torch.long).random_(3).to(device)

        optimizer.zero_grad()
        y_hat = model(x)
        output = loss(y_hat, y)
        output.backward()
        optimizer.step()

fake_train(model, device=device)

quantized_model = convert_fx(model)



</code></pre>
<pre><code class="lang-auto">!pip freeze | grep 'torch'


pytorch-lightning==2.0.3
pytorch-triton==2.1.0+6e4932cda8
torch==2.0.1+cu118
torchaudio==2.0.2+cu118
torchmetrics==0.11.4
torchvision==0.15.2+cu118

</code></pre>

username: Chame_call
cooked: <p><a class="mention" href="/u/jerryzh168">@jerryzh168</a>, any insights?</p>

username: jerryzh168
cooked: <p>sorry for the late reply, this looks pretty weird, I assume it comes from <a href="https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/quantized/cpu/qconv_prepack.cpp" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/quantized/cpu/qconv_prepack.cpp</a> but when I search for “Unsupported”, it looks like per_channel_affine should be captured by other if/else branches</p>

username: jerryzh168
cooked: <p>OK I think it might come from here: <a href="https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/quantized/cudnn/ConvPrepack.cpp#L37" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/quantized/cudnn/ConvPrepack.cpp#L37</a></p>
<p>does it work if you move the model to cpu?</p>

username: Chame_call
cooked: <p>yes, it does. thanks for that but maybe you can have some ideas about the further problem.<br>
I have custom architecture (object detection)  and passing it completely to prepare_fx throws Exceptions due to non-traceable nature some of the operations so I decided to quantize only backbone (classic timm feature extractor) and the only way I found how achieve this is like that:<br>
<code>model.backbone = prepare_fx(model.backbone)</code></p>
<p>after every training epoch I do validation on converted model:</p>
<pre><code class="lang-auto">        quantized_model = copy.deepcopy(model).to('cpu')
        quantized_backbone = convert_fx(quantized_model.backbone)
        quantized_model.backbone = quantized_backbone
</code></pre>
<p>and the performance of the model seems to be very low.<br>
So maybe you have some ideas about the point such as:</p>
<ul>
<li>maybe there’s more preferred way of ignoring submodules from quantization (I tried setting None as qconfig for different modules but that doesn’t work because anyway the whole module that is passed to prepare_fx is going to be traced);</li>
<li>maybe I need to do smth additionally along with my direct backbone substitution;</li>
</ul>
<p>and one more not directly related question:<br>
is there only int8 quantization precision mode available?</p>