I saved the quantized model to TorchScript successfully, but when I run the TorchScript model, I met the following problem:
Traceback (most recent call last):
File "ocr_quant.py", line 293, in <module>
js_out = ts(x)
File "/home/dai/py36env/lib/python3.6/site-packages/torch/nn/modules/module.py", line 541, in __call__
result = self.forward(*input, **kwargs)
RuntimeError: Didn't find kernel to dispatch to for operator 'aten::max_pool2d_with_indices'. Tried to look up kernel for dispatch key 'QuantizedCPUTensorId'. Registered dispatch keys are: [CUDATensorId, CPUTensorId, VariableTensorId]
The above operation failed in interpreter, with the following stack trace:
at <string>:63:30
return torch.avg_pool2d(self, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override), backward
def max_pool2d(self,
kernel_size: List[int],
stride: List[int],
padding: List[int],
dilation: List[int],
ceil_mode: bool):
output, indices = torch.max_pool2d_with_indices(self, kernel_size, stride, padding, dilation, ceil_mode)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
def backward(grad_output):
grad_self = torch.max_pool2d_with_indices_backward(grad_output, self, kernel_size, stride, padding, dilation, ceil_mode, indices)
return grad_self, None, None, None, None, None
return output, backward
def max_pool2d_with_indices(self,
kernel_size: List[int],
stride: List[int],
padding: List[int],
The above operation failed in interpreter, with the following stack trace:
Here’s my code for quantization:
crnn.cnn.fuse_model()
img = Image.open("3.png").convert("L").resize((300,32))
x = ToTensor()(img).unsqueeze(0)
out = crnn(x)
crnn.cnn.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare(crnn.cnn,inplace=True)
# calibration
crnn(x)
# quantize
torch.quantization.convert(crnn.cnn, inplace=True)
torch.quantization.quantize_dynamic(crnn.rnn,{nn.LSTM,nn.Linear},dtype=torch.qint8,inplace=True)
quantized_out = crnn(x)
ts = torch.jit.script(crnn)
ts_out = ts(x)
The error occurs when the TorchScript model ts is called.
And here’s my code for crnn.cnn
class ConvBNReLU(nn.Sequential):
def __init__(self,i,in_channels,out_channels,kernel_size=3,stride = 1,padding = 1,groups = 1,bn = False):
self.bn = bn
if bn:
super(ConvBNReLU,self).__init__(OrderedDict([
("conv{0}".format(i),nn.Conv2d(in_channels,out_channels,kernel_size,stride,padding,groups = groups,bias = True)),
("batchnorm{0}".format(i),nn.BatchNorm2d(out_channels)),
("relu{0}".format(i),nn.ReLU(inplace=False))
]))
else:
super(ConvBNReLU,self).__init__(OrderedDict([
("conv{0}".format(i),nn.Conv2d(in_channels,out_channels,kernel_size,stride,padding,groups = groups,bias = True)),
("relu{0}".format(i),nn.ReLU(inplace=False))
]))
class CNN(nn.Module):
def __init__(self, imgH, nc):
super(CNN, self).__init__()
assert imgH % 16 == 0, 'imgH has to be a multiple of 16'
ks = [3, 3, 3, 3, 3, 3, 2]
ps = [1, 1, 1, 1, 1, 1, 0]
ss = [1, 1, 1, 1, 1, 1, 1]
nm = [64, 128, 256, 256, 512, 512, 512]
self.cnn = nn.Sequential()
pn = 0
for i in range(7):
nIn = nc if i == 0 else nm[i - 1]
nOut = nm[i]
if i in [0,1,3,5]:
self.cnn.add_module("ConvBNReLU{}".format(i),ConvBNReLU(i,nIn,nOut,ks[i],ss[i],ps[i],bn=False))
if i in [0,1]:
self.cnn.add_module("pooling{0}".format(pn),nn.MaxPool2d(2,2))
else:
self.cnn.add_module("pooling{0}".format(pn),nn.MaxPool2d((2,2),(2,1),(0,1)))
pn += 1
else:
self.cnn.add_module("ConvBNReLU{}".format(i),ConvBNReLU(i,nIn,nOut,ks[i],ss[i],ps[i],bn=True))
self.quant = QuantStub()
self.dequant = DeQuantStub()
def forward(self, x):
x = self.quant(x)
x = self.cnn(x)
x = self.dequant(x)
return x
My PyTorch version is 1.3.1+cpu.
How can I solve this problem? Looking forward to your replay.