Implementing a variation of MaxPool2d

vision
1

Hi all, I need to implement a variation of MaxPool2d where all maximum values are selected and their indices are returned in case of equality instead of the first one and adapt MaxUnpool2d to it.

Here is an example,

pool = CustomMaxPool2d(2, stride=2, return_indices=True)
unpool = CustomMaxUnpool2d(2, stride=2)

inp = torch.tensor([[[[ 6.,  2,  3,  4],
                      [ 5,  6,  8,  8],
                      [ 9, 10, 11, 16],
                      [14, 14, 16, 16]]]])

output, indices = pool(inp)
print('maxPool:\n', output)
print('pool indices:\n', indices)
print('maxUnpool:\n', unpool(output, indices))

maxPool:
 tensor([[[[ 6.,  8.],
          [14., 16.]]]])

pool indices:
 tensor([[[[[0, 5], [6, 7]],
          [[12, 13], [11, 14, 15]]]]])

maxUnpool:
 tensor([[[[ 6.,  0.,  0.,  0.],
          [ 0.,  6.,  8.,  8.],
          [ 0.,  0.,  0., 16.],
          [14.,  14.,  16.,  16.]]]])

Is there a way to do this without modifying the c++ code ?

2

I would recommend to start with a manual implementation in Python via unfold to get the patches and apply your custom pooling logic to each patch.
Note that the currently returned indices have a variable shape and would thus not be representable in a standard tensor but e.g. a list instead or a nested tensor (which is an experimental feature at the moment).
Once this approach is working properly, you could try to reimplement it as a custom layer in the C++/CUDA backend.

3

By unfold do you mean torch.nn.unfold or tensor.unfold ?

4

Both should work, but the nn.Unfold module would provide more built-in options such as padding dilation etc.