Torch jit script and FloatFunctional error

AfonsoSalgadoSousa: <p>Hello, I am trying to use the packages of quantization that PyTorch provides to quantize a MobileNet 3D. I followed the steps recommended in <a href="https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html" class="inline-onebox" rel="noopener nofollow ugc">(beta) Static Quantization with Eager Mode in PyTorch — PyTorch Tutorials 2.1.1+cu121 documentation</a>. Everything goes fine until I try to compile the model as a TorchScript with torch.jit.script for saving the model. The “Replacing addition with <code>nn.quantized.FloatFunctional</code>” step yields the following error:</p>
<blockquote>
<p>Traceback (most recent call last):<br>
File “main.py”, line 155, in <br>
main()<br>
File “main.py”, line 146, in main<br>
torch.jit.save(torch.jit.script(fp_model), (opt.model_path / ‘quantized’ / ‘quant_mobilenet3d.pth’).as_posix())<br>
File “/home/ctm/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/jit/<strong>init</strong>.py”, line 1516, in script<br>
return torch.jit._recursive.create_script_module(obj, torch.jit._recursive.infer_methods_to_compile)<br>
File “/home/ctm/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 318, in create_script_module<br>
return create_script_module_impl(nn_module, concrete_type, stubs_fn)<br>
File “/home/ctm/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 372, in create_script_module_impl<br>
script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)<br>
File “/home/ctm/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/jit/<strong>init</strong>.py”, line 1900, in _construct<br>
init_fn(script_module)<br>
File “/home/ctm/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 353, in init_fn<br>
scripted = create_script_module_impl(orig_value, sub_concrete_type, infer_methods_to_compile)<br>
File “/home/ctm/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 372, in create_script_module_impl<br>
script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)<br>
File “/home/ctm/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/jit/<strong>init</strong>.py”, line 1900, in _construct<br>
init_fn(script_module)<br>
File “/home/ctm/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 353, in init_fn<br>
scripted = create_script_module_impl(orig_value, sub_concrete_type, infer_methods_to_compile)<br>
File “/home/ctm/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 376, in create_script_module_impl<br>
create_methods_from_stubs(concrete_type, stubs)<br>
File “/home/ctm/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 292, in create_methods_from_stubs<br>
concrete_type._create_methods(defs, rcbs, defaults)<br>
File “/home/ctm/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/jit/<strong>init</strong>.py”, line 1359, in _recursive_compile_class<br>
_compile_and_register_class(obj, rcb, _qual_name)<br>
File “/home/ctm/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/jit/<strong>init</strong>.py”, line 1363, in _compile_and_register_class<br>
_jit_script_class_compile(qualified_name, ast, rcb)<br>
RuntimeError:<br>
undefined value super:<br>
File “/home/ctm/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/nn/quantized/modules/functional_modules.py”, line 35<br>
def <strong>init</strong>(self):<br>
super(FloatFunctional, self).<strong>init</strong>()<br>
~~~~~ &lt;— HERE<br>
self.activation_post_process = torch.nn.Identity()<br>
‘FloatFunctional.<strong>init</strong>’ is being compiled since it was called from ‘<strong>torch</strong>.torch.nn.quantized.modules.functional_modules.FloatFunctional’<br>
File “/home/ctm/afonso/easyride/acceleration/src/models/mobilenetv2.py”, line 69<br>
if self.use_res_connect:<br>
if self.quantize:<br>
return nn.quantized.FloatFunctional().add(x, self.conv(x))<br>
~~~~~~~~~~~~~~~~~~~~~~~~~~~~ &lt;— HERE<br>
else:<br>
return x + self.conv(x)<br>
‘<strong>torch</strong>.torch.nn.quantized.modules.functional_modules.FloatFunctional’ is being compiled since it was called from ‘InvertedResidual.forward’<br>
File “/home/ctm/afonso/easyride/acceleration/src/models/mobilenetv2.py”, line 69<br>
if self.use_res_connect:<br>
if self.quantize:<br>
return nn.quantized.FloatFunctional().add(x, self.conv(x))<br>
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ &lt;— HERE<br>
else:<br>
return x + self.conv(x)</p>
</blockquote>
<p>Not using the FloatFunctional function and keeping the addition as is enables the model to be saved and loaded correctly, but will later give errors on the inference step because of the operation not being supported by the QuantizedCPU backend. Similarly, using the addition from QFunctional also gives backend-based errors.</p>

AfonsoSalgadoSousa: <p>I found the answer in another thread, but cannot find it again. Basically, the nn.quantized.FloatFunctional() class should be initialized in the <strong>init</strong> function so that it is visible to the TorchScript inspector.<br>
This will not work:</p>
<pre><code class="lang-auto">def forward(self, x):
   return nn.quantized.FloatFunctional().add(x, self.conv(x))
</code></pre>
<p>Yet something like this will:</p>
<pre><code class="lang-auto">def __init__():
   self.q_add = nn.quantized.FloatFunctional()
...
def forward(self, x):
   return self.q_add.add(x, self.conv(x))
</code></pre>