How to do a weighted pooling in Mxnet?

Question

I want to do a 2d convolutional operation that uses same 1x2x4 weight on every channel.

(Note: the input height & width are bigger than our kernel, so I can't just use a dot product.)

How can I do this is mxnet?

I tried to use the same instance of a signle 2d conv layer by concatenating it on every channel, but it is incredibly slow.

def Concat(*args, axis=1, **kwargs):
    net = nn.HybridConcatenate(axis=axis,**kwargs)
    net.add(*args)
    return net

def Seq(*args):
    net = nn.HybridSequential()
    net.add(*args)
    return net

class Trim_D1(nn.HybridBlock):
    def __init__(self, from_, to, **kwargs):
        super(Trim_D1, self).__init__(**kwargs)
        self.from_ = from_
        self.to = to
    def forward(self, x):
        return x[:,self.from_:self.to]


PooPool = nn.Conv2D(kernel_size=(2,4), strides=(2, 4), channels=1, activation=None, use_bias=False, weight_initializer=mx.init.Constant(1/8))

conc = ()
for i in range(40):
    conc += Seq(
                Trim_D1(i,i+1),
                PooPool
            ),

WeightedPool= Concat(*conc)

Ideally I would also want my kernel weights to sum up to 1 in order to resemble the weighted average pooling.

Edit: I think I know how to do this. I'm going to edit Conv2D and _Conv source codes so that instead of creating weights of CxHxW dimension it creates a weight of 1xHxW dimension and uses a broadcasting during the convolutional operation. In order for weights to sum up to 1, additionally a softmax operation has to be applied.

score 0 · Answer 1 · answered Aug 13 '22 at 16:30

Ok, apparently the weights are of in_channels x out_channels x H x W dimensions and broadcasting is not allowed during the convolutional operation. We could fix out_channels to 1 by using the num_groups same as the output channels, as for input channels, we can simply broadcast the same weight n number of times.

In _Conv.__init__ during initialization I discarded the first two dimensions so our kernel is only H x W now:

    self.weight = Parameter('weight', shape=wshapes[1][2:],
                            init=weight_initializer,
                            allow_deferred_init=True)

In _Conv.hybrid_forward I am flattening our weight to 1D in order to perform softmax and then restore to the original 2D shape. Then I expand first two dimensions and repeat the first dimension as mentioned above:

    orig_shape = weight.shape
    act = getattr(F, self._op_name)(x, mx.nd.softmax(weight.reshape(-1)).reshape(orig_shape)[None,None,:].repeat(self._kwargs['num_group'],axis=0), name='fwd', **self._kwargs)

How to do a weighted pooling in Mxnet?

1 Answers1