Operator request - maxpool optional second output

Before I forget, can I request the second optional output of MaxPool to be put on the “to-do” list? (Both as a “layer” output and as an “op” operation e.g. op.MaxPoolIndices) It is not something I need right now, but I have come across it a couple of times in for certain models and I will probably need it in the future. Basically what it does is instead of giving the maximum value when maxpooling, it gives the index of the maximum value.
Here is the pytorch code to create it with the expected results:

import torch
import torch.onnx

class MaxPoolModel(torch.nn.Module):
    def __init__(self):
        super(MaxPoolModel, self).__init__()
        self.maxpool = torch.nn.MaxPool2d(kernel_size=(2, 2), stride=1, return_indices=True)

    def forward(self, x):
        # Output the maxpool result
        _, indices = self.maxpool(x)
        return indices

# Create a model instance
model = MaxPoolModel()

model.eval()

# Create a random input tensor
input = torch.randn(1, 1, 4, 4)

# Output the model's predictions
output = model(input)

print(input)

print(output)

#export the model
torch.onnx.export(model,           
                  input,                  
                  "maxpool_model.onnx",  
                  export_params=True,      
                  opset_version=12,        
                  do_constant_folding=True,
                  input_names = ['input'],  
                  output_names = ['output'],
                  dynamic_axes={'input' : {0 : 'batch_size'},  
                                'output' : {0 : 'batch_size'}})

yup good idea :slight_smile:

1 Like

Thank you for reporting this! It’s now known internally as task 179.

1 Like