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?