QAT: trace_model 'QuantizedCPU' backend error

ynjiun_wang: <p>I have an QAT quantized_model which has no problem to run:</p>
<pre><code class="lang-auto">quantized_model.eval()
_ = quantized_model(torch.rand(1,3,300,300))
</code></pre>
<p>and it also can be traced successfuly:</p>
<pre><code class="lang-auto">trace_model = torch.jit.trace(quantized_model, torch.Tensor(1,3,300,300))
</code></pre>
<p>but when I tried to run the trace_model as below:</p>
<pre><code class="lang-auto">trace_model.eval()
 _ = trace_model(torch.rand(1,3,300,300))
</code></pre>
<p>I encountered the following error message:</p>
<pre><code class="lang-auto">  _ = trace_model(torch.rand(1,3,300,300))
  File "/home/paul/rknn2/lib/python3.6/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "&lt;string&gt;", line 63, in &lt;foward op&gt;
                       dilation: List[int],
                       ceil_mode: bool):
            output, indices = torch.max_pool2d_with_indices(self, kernel_size, stride, padding, dilation, ceil_mode)
                              ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ &lt;--- HERE
            def backward(grad_output):
                grad_self = torch.max_pool2d_with_indices_backward(grad_output, self, kernel_size, stride, padding, dilation, ceil_mode, indices)
RuntimeError: Could not run 'aten::max_pool2d_with_indices' with arguments from the 'QuantizedCPU' backend. 'aten::max_pool2d_with_indices' is only available for these backends: [CPU, CUDA, Named, Autograd, Profiler, Tracer].
</code></pre>
<p>What am I missing? How can I fix this error? Thank you for your help.</p>

HDCharles: <p>looks like the quantized torchscript is not playing well with the autograd, try running without gradient</p>
<p>with torch.no_grad():<br>
_ = trace_model(torch.rand(1,3,300,300))</p>
<p>otherwise not sure, if you can give a minimal repro that would help.</p>

ynjiun_wang: <p><a class="mention" href="/u/hdcharles">@HDCharles</a></p>
<p>Thanks for the help. It works now!</p>
<p>The reason I create such test script is to duplicate the same “QuantizedCPU” backend error I encountered when I saved the trace_model into a file, says, “trace_model.pt” and then tried to load it by using</p>
<pre><code class="lang-auto">rknn.load_pytorch(model="trace_model.pt", input_size_list=input_size_list)
</code></pre>
<p>where rknn is a toolkit for RockChip NPU.</p>
<p>Now my question is: is there a way to save the trace_model with torch.no_grad() ? such that when rknn.load_pytorch would not encounter “QuantizedCPU” backend error?</p>
<p>or has to modify rknn.load_pytorch module to add</p>
<p><code>with torch.no_grad():</code></p>
<p>in it before parsing the model script?</p>
<p>Thank you again for your help.</p>

HDCharles: <p>you can try just shoving that into the forward loop of the model, not sure if that would work.</p>
<p>note: dont do something like:</p>
<p>if train_flag:<br>
code<br>
else:<br>
with torch.no_grad<br>
code</p>
<p>because the if statement will break the tracing.</p>
<p>alternatively you can try:</p>
<p>with torch.no_grad()<br>
rknn.load_pytorch(…)</p>