Maxpool2d doesn't work in TorchScript after quantization?

dalalaa: <p>I saved the quantized model to TorchScript successfully, but when I run the TorchScript model, I met the following problem:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "ocr_quant.py", line 293, in  &lt;module&gt;
    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 &lt;string&gt;: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)
                              ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ &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)
                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:
</code></pre>
<p>Here’s my code for quantization:</p>
<pre><code class="lang-auto">    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)
</code></pre>
<p>The error occurs when the TorchScript model <code>ts</code> is called.</p>
<p>And here’s my code for crnn.cnn</p>
<pre><code class="lang-auto">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
</code></pre>
<p>My PyTorch version is 1.3.1+cpu.</p>
<p>How can I solve this problem? Looking forward to your replay.</p>

dalalaa: <p>The input tensor of CRNN is cropped from an tensor of image, when I convert the tensor to PIL image and convert it back, the code works well.<br>
<code>img_out = ToTensor()(ToPILImage()(img_out.squeeze(0))).unsqueeze(0)</code></p>
<p>and the following code also helps:</p>
<pre><code class="lang-auto">img_out = torch.ones(img_out_.shape)
img_out.data = img_out_.clone()
</code></pre>
<p>Could anyone tell me the differences between cropped tensor and tensor converted from image?</p>

eellison: <p>Hi <a class="mention" href="/u/dalalaa">@dalalaa</a>, that sounds like a bug. Do you mind raising an issue on github ?</p>

dalalaa: <p>Thank you for your reply, I will try to find an minimum code to reproduce it.</p>

ahs8w: <p>I know this is late but for others visiting this with the same error I found that wrapping the <code>ts(x)</code> in a <code>torch.no_grad()</code> block solved the error</p>