I’m now implementing a pooling layer similar to ‘super-pixel Pooling’ 89 which has pre-computed superpixel masks to guide the pooling.
Firstly I read the document about extending pytorch 16 which says

You can extend it in both ways, but we recommend using modules for all kinds of layers, that hold any parameters or buffers, and recommend using a functional form parameter-less operations like activation functions, pooling, etc.

So I got confused here, because the operator I want to implement has no parameter, but do have a buffer (the superpixel masks to guide the pooling).
Now I just finished my implementation based on ‘Module’, by inheriting the nn.Module and define the forward():
class SuperPixelPooling(nn.Module):
  def __init__(self. input, mask, output):
    super(SuperPixelPooling, self).__init__()
  def forward(self, input, mask)
    # perform superpixel pooling on input according to mask
    output = do_superpixel_pooling(input, mask)
    return output

All the operations in forward() are based on Variables.
Another problem occurs: when I’m using the new ‘SuperpixelPooling’ module in my model, it seems impossible to put such module with multiple inputs into my model.
class LeNet(nn.Module):
  def __init__(self):
    super(LeNet, self).__init__()
    self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
    self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
    # seems not right?##############
    self.sp_pool2 = SuperPixelPooling()
    # ##############################
    self.conv2_drop = nn.Dropout2d()
    self.fc1 = nn.Linear(320, 50)
    self.fc2 = nn.Linear(50, 10)

  def forward(self, x):
    x = F.relu(F.max_pool2d(self.conv1(x), 2))
    x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
    # or like this ? #################
    x = SuperPixelPooling(x, mask)
    ##################################
    x = x.view(-1, 320)
    x = F.relu(self.fc1(x))
    x = F.dropout(x, training=self.training)
    x = self.fc2(x)
    return F.log_softmax(x)

I have to use this module with multiple inputs as a Function in my model’s forward(), right?