Can't use pytorch operator: max_pool2d_with_indices in coremltools

Mobile
1

Nowadays, a lot of PyTorch models use MaxPool2d operator with the option return_indices=True

I need to convert some PyTorch models into CoreML. For that I use coremltools. But coremltools doesn’t support yet this operator with return_indices=True. And it is confirmed in a GitHub issue on their repo

When you try using this code:

traced_model = torch.jit.trace(model, input_t)
mlmodel = ct.convert(
    traced_model,
    inputs=[ct.TensorType(shape=input_t.shape)]
)

You’ve got this runtime error:

RuntimeError: PyTorch convert function for op 'max_pool2d_with_indices' not implemented

After reading the doc of coremltools, you’ve got 2 choices

  • Implementing a custom operator in coremltools: Not as easy as it seems
  • Decompose the PyTorch model causing the issue into coremltools supported operators: That’s what I try to do, but I’m debutant in PyTorch.

So I need help in order to decompose this:

pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, return_indices=True)
output, indices = pool(input)

Into something like

pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, return_indices=False) # <-- notice the return_indices=False here
output = pool(input)
indices = some code here...

Please could someone help me?