Pytorch network quntization aware training problem

Shixian_Wen: <p>Hi:</p>
<p>I am trying to use quantization aware training for my CNN network. I used the lastest 1.90 pytorch nightly version. What I want to do is I load a pretrained RestNet18 and finetune it with other dataset. I run into some problem to quantitize the network during training.</p>
<p>I followed the <a href="https://pytorch.org/tutorials/intermediate/quantized_transfer_learning_tutorial.html" rel="noopener nofollow ugc">tutorial</a> .</p>
<p>However, I run into an NotImplementedError: “Could not run ‘aten::thnn_conv2d_forward’ with arguments from the ‘QuantizedCPU’ backend.” when I want to use the converted int8 model for testing. In addition, I found the saved converted int8 model is not quantitized properly, it is still represented as floating points number.</p>
<p>Could you help me understand this error?</p>
<p>Many thanks<br>
Shixian Wen</p>
<p><strong>Error meesage</strong></p>
<p>Traceback (most recent call last):<br>
File “attetion_delta_quantization_tutorial.py”, line 902, in <br>
acc = test_cpu(testLoader_mix10,model_one_int8)<br>
File “attetion_delta_quantization_tutorial.py”, line 708, in test_cpu<br>
output = model(data)<br>
File “/home/shixian/anaconda3/envs/pytorchtest/lib/python3.6/site-packages/torch/nn/modules/module.py”, line 1015, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “attetion_delta_quantization_tutorial.py”, line 552, in forward<br>
x = self.model_fe_features(x)<br>
File “/home/shixian/anaconda3/envs/pytorchtest/lib/python3.6/site-packages/torch/nn/modules/module.py”, line 1015, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/home/shixian/anaconda3/envs/pytorchtest/lib/python3.6/site-packages/torch/nn/modules/container.py”, line 139, in forward<br>
input = module(input)<br>
File “/home/shixian/anaconda3/envs/pytorchtest/lib/python3.6/site-packages/torch/nn/modules/module.py”, line 1015, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/home/shixian/anaconda3/envs/pytorchtest/lib/python3.6/site-packages/torch/nn/modules/container.py”, line 139, in forward<br>
input = module(input)<br>
File “/home/shixian/anaconda3/envs/pytorchtest/lib/python3.6/site-packages/torch/nn/modules/module.py”, line 1015, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/home/shixian/anaconda3/envs/pytorchtest/lib/python3.6/site-packages/torch/nn/modules/conv.py”, line 433, in forward<br>
return self._conv_forward(input, self.weight, self.bias)<br>
File “/home/shixian/anaconda3/envs/pytorchtest/lib/python3.6/site-packages/torch/nn/modules/conv.py”, line 430, in _conv_forward<br>
self.padding, self.dilation, self.groups)<br>
NotImplementedError: Could not run ‘aten::thnn_conv2d_forward’ with arguments from the ‘QuantizedCPU’ 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::thnn_conv2d_forward’ is only available for these backends: [CPU, CUDA, BackendSelect, Named, InplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, UNKNOWN_TENSOR_TYPE_ID, AutogradMLC, AutogradNestedTensor, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, Tracer, Autocast, Batched, VmapMode].</p>
<p>CPU: registered at /pytorch/build/aten/src/ATen/RegisterCPU.cpp:11727 [kernel]<br>
CUDA: registered at /pytorch/build/aten/src/ATen/RegisterCUDA.cpp:13375 [kernel]<br>
BackendSelect: fallthrough registered at /pytorch/aten/src/ATen/core/BackendSelectFallbackKernel.cpp:3 [backend fallback]<br>
Named: registered at /pytorch/aten/src/ATen/core/NamedRegistrations.cpp:7 [backend fallback]<br>
InplaceOrView: fallthrough registered at /pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:60 [backend fallback]<br>
AutogradOther: registered at /pytorch/torch/csrc/autograd/generated/VariableType_0.cpp:9539 [autograd kernel]<br>
AutogradCPU: registered at /pytorch/torch/csrc/autograd/generated/VariableType_0.cpp:9539 [autograd kernel]<br>
AutogradCUDA: registered at /pytorch/torch/csrc/autograd/generated/VariableType_0.cpp:9539 [autograd kernel]<br>
AutogradXLA: registered at /pytorch/torch/csrc/autograd/generated/VariableType_0.cpp:9539 [autograd kernel]<br>
UNKNOWN_TENSOR_TYPE_ID: registered at /pytorch/torch/csrc/autograd/generated/VariableType_0.cpp:9539 [autograd kernel]<br>
AutogradMLC: registered at /pytorch/torch/csrc/autograd/generated/VariableType_0.cpp:9539 [autograd kernel]<br>
AutogradNestedTensor: registered at /pytorch/torch/csrc/autograd/generated/VariableType_0.cpp:9539 [autograd kernel]<br>
AutogradPrivateUse1: registered at /pytorch/torch/csrc/autograd/generated/VariableType_0.cpp:9539 [autograd kernel]<br>
AutogradPrivateUse2: registered at /pytorch/torch/csrc/autograd/generated/VariableType_0.cpp:9539 [autograd kernel]<br>
AutogradPrivateUse3: registered at /pytorch/torch/csrc/autograd/generated/VariableType_0.cpp:9539 [autograd kernel]<br>
Tracer: registered at /pytorch/torch/csrc/autograd/generated/TraceType_0.cpp:9373 [kernel]<br>
Autocast: fallthrough registered at /pytorch/aten/src/ATen/autocast_mode.cpp:250 [backend fallback]<br>
Batched: registered at /pytorch/aten/src/ATen/BatchingRegistrations.cpp:1016 [backend fallback]<br>
VmapMode: fallthrough registered at /pytorch/aten/src/ATen/VmapModeRegistrations.cpp:33 [backend fallback]</p>
<p><strong>Here is my model:</strong></p>
<p>model_fe = models.resnet18(pretrained=True, progress=True, quantize=False)<br>
num_ftrs = model_fe.fc.in_features<br>
model_fe.train()<br>
model_fe.fuse_model()</p>
<p>class Resnet18_ONE(nn.Module):</p>
<pre><code>def __init__(self,model_fe=model_fe):

    super(Resnet18_ONE,self).__init__()
    #self.loss = loss
    self.model_fe_features = nn.Sequential(model_fe.quant,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)
    #print(self.base)
    self.linear_sub = nn.Linear(num_ftrs, superclasses)
    self.linear_bird = nn.Linear(num_ftrs, classes_bird)
    self.linear_boat = nn.Linear(num_ftrs, classes_boat)
    self.linear_car = nn.Linear(num_ftrs, classes_car)
    self.linear_cat = nn.Linear(num_ftrs, classes_cat)
    self.linear_fungus = nn.Linear(num_ftrs, classes_fungus)
    self.linear_insect = nn.Linear(num_ftrs, classes_insect)
    self.linear_monkey = nn.Linear(num_ftrs, classes_monkey)
    self.linear_truck = nn.Linear(num_ftrs, classes_truck)
    self.linear_dog = nn.Linear(num_ftrs, classes_dog)
    self.linear_fruit = nn.Linear(num_ftrs, classes_fruit)
    
def forward(self,x):

    x = self.model_fe_features(x)
    x = torch.flatten(x, 1)
    if task == 'SUB':
        #print("I am in sub")
        x = self.linear_sub(x)
    elif task == 'BIRD':
        #print("I am in bird")
        x = self.linear_bird(x)
    elif task == 'BOAT':
        x = self.linear_boat(x)
    elif task == 'CAR':
        x = self.linear_car(x)
    elif task == 'CAT':
        x = self.linear_cat(x)
    elif task == 'FUNGUS':
        x = self.linear_fungus(x)
    elif task == 'INSECT':
        x = self.linear_insect(x)
    elif task == 'MONKEY':
        x = self.linear_monkey(x)
    elif task == 'TRUCK':
        x = self.linear_truck(x)
    elif task == 'DOG':
        x = self.linear_dog(x)
    else:
        #print("enter fruit")
        x = self.linear_fruit(x)
               
    return x
</code></pre>
<p><strong>Here is how I set and train the quantization model:</strong></p>
<p>model_one = Resnet18_ONE(model_fe=model_fe)<br>
model_one.qconfig = torch.quantization.get_default_qat_qconfig(‘fbgemm’)<br>
model_one_prepared = torch.quantization.prepare(model_one,inplace=True)<br>
optimizer_sub = optim.SGD(filter(lambda p: p.requires_grad, model_one_prepared.parameters()), lr=args.lr, momentum=args.momentum)<br>
train(epoch,trainLoader_mix10,optimizer_sub,model_one_prepared)</p>
<p><strong>After training the model, Here is how I save and covert the model to int8:</strong></p>
<p>model_one_prepared.to(‘cpu’)<br>
model_one_int8 = torch.quantization.convert(model_one_prepared.eval(),inplace=False)<br>
model_one_prepared.eval()<br>
model_one_int8.eval()<br>
acc = test_cpu(testLoader_mix10,model_one_int8)<br>
torch.save({‘modelONE_state_dict’: model_one_int8.state_dict(),}, save_path_quantize)</p>