Can QAT inference on CUDA?

username: Heaseo_Chung
cooked: <p>Hi,</p>
<p>I know that static &amp; dynamic quantization cannot inference with CUDA.</p>
<p>but I am wondering that QAT model can inference with CUDA.</p>
<p>Thanks.</p>

username: rsundar96
cooked: <p>From the <a href="https://pytorch.org/docs/stable/quantization.html" rel="noopener nofollow ugc">PyTorch Quantization docs</a></p>
<p><code>Quantization-aware training (through FakeQuantize) supports both CPU and CUDA</code></p>

username: Heaseo_Chung
cooked: <p>Quantization doc says that it does support both CPU and GPU. I tried the tutorial and it didn’t work.<br>
I am still confusing because some of users are saying it does not support GPU yet from <a href="https://discuss.pytorch.org/t/is-qat-inference-not-support-gpu/109901" class="inline-onebox">Is QAT Inference not support GPU?</a></p>

username: raghuramank100
cooked: <p>You can run a QAT model prior to convert on GPU. Please look at the example in torchvision: <a href="https://github.com/pytorch/vision/blob/master/references/classification/train_quantization.py#L130" class="inline-onebox" rel="noopener nofollow ugc">vision/train_quantization.py at master · pytorch/vision · GitHub</a></p>

username: Heaseo_Chung
cooked: <h1>
<a name="after-i-applied-qat-method-and-i-tried-to-inference-the-model-with-gpu-but-i-got-this-error-below-however-cpu-is-working-fine" class="anchor" href="#after-i-applied-qat-method-and-i-tried-to-inference-the-model-with-gpu-but-i-got-this-error-below-however-cpu-is-working-fine"></a>After I applied QAT method and I tried to inference the model with GPU but I got this error below. However CPU is working fine.</h1>
<p>File “quantize_model.py”, line 359, in <br>
model = quantization_aware_training(model, device)<br>
File “quantize_model.py”, line 120, in quantization_aware_training<br>
torch.quantization.convert(quantized_eval_model, inplace=True)<br>
File “/opt/conda/lib/python3.8/site-packages/torch/quantization/quantize.py”, line 471, in convert<br>
_convert(<br>
File “/opt/conda/lib/python3.8/site-packages/torch/quantization/quantize.py”, line 507, in _convert<br>
_convert(mod, mapping, True,  # inplace<br>
File “/opt/conda/lib/python3.8/site-packages/torch/quantization/quantize.py”, line 509, in _convert<br>
reassign[name] = swap_module(mod, mapping, custom_module_class_mapping)<br>
File “/opt/conda/lib/python3.8/site-packages/torch/quantization/quantize.py”, line 534, in swap_module<br>
new_mod = mapping[type(mod)].from_float(mod)<br>
File “/opt/conda/lib/python3.8/site-packages/torch/nn/intrinsic/quantized/modules/conv_relu.py”, line 97, in from_float<br>
return super(ConvReLU2d, cls).from_float(mod)<br>
File “/opt/conda/lib/python3.8/site-packages/torch/nn/quantized/modules/conv.py”, line 418, in from_float<br>
return _ConvNd.from_float(cls, mod)<br>
File “/opt/conda/lib/python3.8/site-packages/torch/nn/quantized/modules/conv.py”, line 220, in from_float<br>
return cls.get_qconv(mod, activation_post_process, weight_post_process)<br>
File “/opt/conda/lib/python3.8/site-packages/torch/nn/quantized/modules/conv.py”, line 187, in get_qconv<br>
qweight = _quantize_weight(mod.weight.float(), weight_post_process)<br>
File “/opt/conda/lib/python3.8/site-packages/torch/nn/quantized/modules/utils.py”, line 14, in _quantize_weight<br>
qweight = torch.quantize_per_channel(<br>
RuntimeError: Could not run ‘aten::quantize_per_channel’ with arguments from the ‘CUDA’ backend. This could be because the operator doesn’t exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit <a href="https://fburl.com/ptmfixes" rel="noopener nofollow ugc">https://fburl.com/ptmfixes</a> for possible resolutions. ‘aten::quantize_per_channel’ is only available for these backends: [CPU, BackendSelect, Named, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, AutogradNestedTensor, UNKNOWN_TENSOR_TYPE_ID, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, Tracer, Autocast, Batched, VmapMode].</p>
<p>Quantization<br>
model.to(device)<br>
model.eval()<br>
model.fuse_model()<br>
model.qconfig = torch.quantization.get_default_qat_qconfig(‘fbgemm’)<br>
torch.quantization.prepare_qat(model, inplace=True)</p>
<p>Evaluation<br>
with torch.no_grad():<br>
model.eval()<br>
epoch_psnr = AverageMeter()<br>
quantized_eval_model = copy.deepcopy(model_without_ddp)<br>
quantized_eval_model.eval()<br>
quantized_eval_model.to(device)<br>
torch.quantization.convert(quantized_eval_model, inplace=True)</p>
<pre><code>        for data in eval_dataloader:
            inputs, labels = data

            inputs = inputs.to(device, non_blocking=True)
            labels = labels.to(device, non_blocking=True)
            preds = quantized_eval_model(inputs).clamp(0.0, 1.0)
</code></pre>
<p>Model<br>
self.quant = torch.quantization.QuantStub()<br>
self.conv_relu1 = ConvReLu(1, 64, _kernel_size=5, _padding=5//2)<br>
self.conv_relu2 = ConvReLu(64,32, _kernel_size=3, _padding=3//2)<br>
self.sub_pixel = nn.Sequential(<br>
nn.Conv2d(32, 1 * (scale_factor ** 2), kernel_size=3, stride=1, padding=1),<br>
nn.PixelShuffle(scale_factor),<br>
nn.Sigmoid()<br>
)<br>
self.dequant = torch.quantization.DeQuantStub()</p>
<pre><code>def forward(self, x):
    x = self.quant(x)
    x = self.conv_relu1(x)
    x = self.conv_relu2(x)
    x = self.sub_pixel(x)
    x = self.dequant(x)
    return x
</code></pre>
<p>class ConvReLu(nn.Sequential):<br>
def <strong>init</strong>(self, _in_channels, _out_channels, _kernel_size, _padding=0):<br>
super(ConvReLu, self).<strong>init</strong>(<br>
nn.Conv2d(_in_channels, _out_channels, kernel_size=_kernel_size, padding=_padding),<br>
nn.ReLU(_out_channels)<br>
)</p>