Unable to convert a quantized model to TorchScript

username: devesh-k
cooked: <p>Hi,<br>
I am following the official tutorials <a href="https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html" rel="noopener nofollow ugc">here</a> and <a href="https://pytorch.org/tutorials/intermediate/quantized_transfer_learning_tutorial.html" rel="noopener nofollow ugc">here</a> to quantize a model but it is errors out while saving to TorchScript.</p>
<p>Below is the code to reproduce this error:<br>
Step 1 - imports</p>
<pre><code class="lang-auto">import timm
import torch
import torch.nn as nn
import torchaudio.transforms as AT
import torchvision.transforms as VT
from nnAudio import features
</code></pre>
<p>Step 2 : Define methods as per the <a href="https://pytorch.org/tutorials/intermediate/quantized_transfer_learning_tutorial.html" rel="noopener nofollow ugc">tutorial</a></p>
<pre><code class="lang-auto">def quant_prep(model_ft):
    model_ft[0].qconfig = torch.quantization.default_qat_qconfig  # Use default QAT configuration
# Step 3
    model_ft = torch.quantization.prepare_qat(model_ft, inplace=True)
    for param in model_ft.parameters():
        param.requires_grad = True
    return(model_ft)
</code></pre>
<pre><code class="lang-auto">import torchvision.models.quantization as models

model = models.resnet18(pretrained=True, progress=True, quantize=False)
num_ftrs = model.fc.in_features
model.train()
model.fuse_model()
</code></pre>
<pre><code class="lang-auto">from torch import nn

def create_combined_model(model_fe):
    # Step 1. Isolate the feature extractor.
    model_fe_features = nn.Sequential(
    model_fe.quant,  # Quantize the input
    model_fe.conv1,
    model_fe.bn1,
    model_fe.relu,
    model_fe.maxpool,
    model_fe.layer1,
    model_fe.layer2,
    model_fe.layer3,
    model_fe.layer4,
    model_fe.avgpool,
    model_fe.dequant,  # Dequantize the output
  )

    # Step 2. Create a new "head"
    new_head = nn.Sequential(
    nn.Dropout(p=0.5),
    nn.Linear(num_ftrs, 2),
   )

  # Step 3. Combine, and don't forget the quant stubs.
    new_model = nn.Sequential(
    model_fe_features,
    nn.Flatten(1),
    new_head,)
    
    model_ft  = quant_prep(new_model)
    
    return model_ft
</code></pre>
<p>Step 3 - define custom “model”</p>
<pre><code class="lang-auto">class MyModel(nn.Module):
    def __init__(self,model = model,image_size = 224):
        super().__init__()
        # num_classes=0 removes the pretrained head
        self.backbone = create_combined_model(model)
        self.sizer = VT.Resize((image_size,image_size),antialias = True)
        self.spec_layer = AT.Spectrogram(n_fft = int(config.NFFT), return_complex= False,).to('cuda')
        self.batch_norm = nn.BatchNorm2d(num_features= 1)
               
        
    def forward(self, x,train = True):
        # first compute spectrogram
        spec_gram = self.spec_layer(x).to('cuda')
        spec_gram = self.batch_norm(spec_gram)
        spec_gram_nan_check = torch.isnan(spec_gram).any().item()
        assert not (spec_gram_nan_check) ,"Tensor contains NaN values after spec gram creation."
        
        with torch.no_grad():
            if train == True:
                #generate a random number and if condition is met apply aug
                ta_transformations_rndm_choice = VT.RandomChoice([AT.FrequencyMasking(freq_mask_param=100),AT.TimeMasking(time_mask_param=50)], p=[.4, .4])
                ta_transformations_rndm_apply = VT.RandomApply([AT.FrequencyMasking(freq_mask_param=50),AT.TimeMasking(time_mask_param=25)],p = .15)
                spec_gram = ta_transformations_rndm_choice(spec_gram)
                spec_gram = ta_transformations_rndm_apply(spec_gram)
                spec_gram_nan_check = torch.isnan(spec_gram).any().item()
                assert not (spec_gram_nan_check) ,"Tensor contains NaN values after augmentations  "
                
                    
        
        x = self.sizer(spec_gram.squeeze(dim = 1))
        x = x.unsqueeze(dim = 1)
              
        # then repeat channels
        del spec_gram,spec_gram_nan_check
        if DEBUG:
            print("Final shape that goes to backbone = " + str(x.shape))
                
        x = x.expand(-1, 3, -1, -1)
        x = self.backbone(x)
        backbone_op_nan_check = torch.isnan(x).any().item()
        assert not (backbone_op_nan_check) ,"Tensor contains NaN values in the backbone OP "
        pred = x
        output = {"prediction": pred }
        #print(output)
        del x , backbone_op_nan_check
        return output

</code></pre>
<p>Step 4 - create an instance of MyModel</p>
<pre><code class="lang-auto">model_ft_tuned_new = MyModel(model)
</code></pre>
<p>Step 5:</p>
<pre><code class="lang-auto">from torch.quantization import convert
model_ft_tuned_new.cpu()
model_quantized_and_trained = convert(model_ft_tuned_new, inplace=False)
</code></pre>
<p>Step 6:- converting to torchscript( this is where the model errors out):</p>
<pre><code class="lang-auto">torch.jit.script(model_quantized_and_trained ,"resnet_ft.pt")
</code></pre>
<p>The stack trace of the error is as below :</p>
<pre><code class="lang-auto">/opt/conda/lib/python3.8/site-packages/torch/jit/_script.py:1277: UserWarning: `optimize` is deprecated and has no effect. Use `with torch.jit.optimized_execution() instead
  warnings.warn(
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
/tmp/ipykernel_1270/2629330692.py in &lt;module&gt;
----&gt; 1 torch.jit.script(model_quantized_and_trained ,"resnet_ft.pt")

/opt/conda/lib/python3.8/site-packages/torch/jit/_script.py in script(obj, optimize, _frames_up, _rcb, example_inputs)
   1322     if isinstance(obj, torch.nn.Module):
   1323         obj = call_prepare_scriptable_func(obj)
-&gt; 1324         return torch.jit._recursive.create_script_module(
   1325             obj, torch.jit._recursive.infer_methods_to_compile
   1326         )

/opt/conda/lib/python3.8/site-packages/torch/jit/_recursive.py in create_script_module(nn_module, stubs_fn, share_types, is_tracing)
    557     if not is_tracing:
    558         AttributeTypeIsSupportedChecker().check(nn_module)
--&gt; 559     return create_script_module_impl(nn_module, concrete_type, stubs_fn)
    560 
    561 

/opt/conda/lib/python3.8/site-packages/torch/jit/_recursive.py in create_script_module_impl(nn_module, concrete_type, stubs_fn)
    634     # Compile methods if necessary
    635     if concrete_type not in concrete_type_store.methods_compiled:
--&gt; 636         create_methods_and_properties_from_stubs(
    637             concrete_type, method_stubs, property_stubs
    638         )

/opt/conda/lib/python3.8/site-packages/torch/jit/_recursive.py in create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
    467     property_rcbs = [p.resolution_callback for p in property_stubs]
    468 
--&gt; 469     concrete_type._create_methods_and_properties(
    470         property_defs, property_rcbs, method_defs, method_rcbs, method_defaults
    471     )

RuntimeError: 
Arguments for call are not valid.
The following variants are available:
  
  aten::eq.Tensor(Tensor self, Tensor other) -&gt; Tensor:
  Expected a value of type 'Tensor' for argument 'other' but instead found type 'bool'.
  
  aten::eq.Scalar(Tensor self, Scalar other) -&gt; Tensor:
  Expected a value of type 'number' for argument 'other' but instead found type 'bool'.
  
  aten::eq.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -&gt; Tensor(a!):
  Expected a value of type 'number' for argument 'other' but instead found type 'bool'.
  
  aten::eq.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -&gt; Tensor(a!):
  Expected a value of type 'Tensor' for argument 'other' but instead found type 'bool'.
  
  aten::eq.int_list(int[] a, int[] b) -&gt; bool:
  Expected a value of type 'List[int]' for argument 'a' but instead found type 'Tensor (inferred)'.
  Inferred the value for argument 'a' to be of type 'Tensor' because it was not annotated with an explicit type.
  
  aten::eq.device(Device a, Device b) -&gt; bool:
  Expected a value of type 'Device' for argument 'a' but instead found type 'Tensor (inferred)'.
  Inferred the value for argument 'a' to be of type 'Tensor' because it was not annotated with an explicit type.
  
  aten::eq.bool(bool a, bool b) -&gt; bool:
  Expected a value of type 'bool' for argument 'a' but instead found type 'Tensor (inferred)'.
  Inferred the value for argument 'a' to be of type 'Tensor' because it was not annotated with an explicit type.
  
  aten::eq.enum(AnyEnumType a, AnyEnumType b) -&gt; bool:
  Expected a value of type 'AnyEnumType' for argument 'a' but instead found type 'Tensor (inferred)'.
  Inferred the value for argument 'a' to be of type 'Tensor' because it was not annotated with an explicit type.
  
  aten::eq.int(int a, int b) -&gt; bool:
  Expected a value of type 'int' for argument 'b' but instead found type 'bool'.
  
  aten::eq.complex(complex a, complex b) -&gt; bool:
  Expected a value of type 'complex' for argument 'b' but instead found type 'bool'.
  
  aten::eq.float(float a, float b) -&gt; bool:
  Expected a value of type 'float' for argument 'b' but instead found type 'bool'.
  
  aten::eq.int_float(int a, float b) -&gt; bool:
  Expected a value of type 'float' for argument 'b' but instead found type 'bool'.
  
  aten::eq.float_int(float a, int b) -&gt; bool:
  Expected a value of type 'int' for argument 'b' but instead found type 'bool'.
  
  aten::eq.float_complex(float a, complex b) -&gt; bool:
  Expected a value of type 'complex' for argument 'b' but instead found type 'bool'.
  
  aten::eq.complex_float(complex a, float b) -&gt; bool:
  Expected a value of type 'float' for argument 'b' but instead found type 'bool'.
  
  aten::eq(Scalar a, Scalar b) -&gt; bool:
  Expected a value of type 'number' for argument 'b' but instead found type 'bool'.
  
  aten::eq.str(str a, str b) -&gt; bool:
  Expected a value of type 'str' for argument 'a' but instead found type 'Tensor (inferred)'.
  Inferred the value for argument 'a' to be of type 'Tensor' because it was not annotated with an explicit type.
  
  aten::eq.float_list(float[] a, float[] b) -&gt; bool:
  Expected a value of type 'List[float]' for argument 'a' but instead found type 'Tensor (inferred)'.
  Inferred the value for argument 'a' to be of type 'Tensor' because it was not annotated with an explicit type.
  
  aten::eq.Tensor_list(Tensor[] a, Tensor[] b) -&gt; bool:
  Expected a value of type 'List[Tensor]' for argument 'a' but instead found type 'Tensor (inferred)'.
  Inferred the value for argument 'a' to be of type 'Tensor' because it was not annotated with an explicit type.
  
  aten::eq.bool_list(bool[] a, bool[] b) -&gt; bool:
  Expected a value of type 'List[bool]' for argument 'a' but instead found type 'Tensor (inferred)'.
  Inferred the value for argument 'a' to be of type 'Tensor' because it was not annotated with an explicit type.
  
  aten::eq.str_list(str[] a, str[] b) -&gt; bool:
  Expected a value of type 'List[str]' for argument 'a' but instead found type 'Tensor (inferred)'.
  Inferred the value for argument 'a' to be of type 'Tensor' because it was not annotated with an explicit type.
  
  eq(float a, Tensor b) -&gt; Tensor:
  Expected a value of type 'Tensor' for argument 'b' but instead found type 'bool'.
  
  eq(int a, Tensor b) -&gt; Tensor:
  Expected a value of type 'Tensor' for argument 'b' but instead found type 'bool'.
  
  eq(complex a, Tensor b) -&gt; Tensor:
  Expected a value of type 'Tensor' for argument 'b' but instead found type 'bool'.

The original call is:
  File "/tmp/ipykernel_1270/2151903598.py", line 33
        
        with torch.no_grad():
            if train == True:
               ~~~~~~~~~~~~~ &lt;--- HERE
                #generate a random number and if condition is met apply aug
                ta_transformations_rndm_choice = VT.RandomChoice([AT.FrequencyMasking(freq_mask_param=100),AT.TimeMasking(time_mask_param=50)], p=[.4, .4])
</code></pre>

username: jerryzh168
cooked: <p>it’s probably because this is just not scriptable, could you remove the <code>if train == True</code> code?</p>

username: devesh-k
cooked: <p>Thanks Jerry. It started working after removing the <code>if train == True </code> part.<br>
Appreciate your help.</p>