How to implement SpatialDropout2D layer?

Gluon
1

Hi everyone, I’m new to both deep learning and MxNet. Can you help me to implement SpatialDropout2D from Keras in MxNet?

2

It’s fairly simple to do in mxnet. Here is a simple gluon Block that implements SpatialDroput2D:

class SpatialDropout2D(gluon.Block):
    def __init__(self, p):
        super(SpatialDropout2D, self).__init__()
        self.p = p

    def forward(self, x):
        if not autograd.is_training():
            return x
        mask_shape = x.shape[:2] + (1, 1)
        mask = nd.random.multinomial(nd.array([self.p, 1 - self.p]),
                                     shape=mask_shape).astype('float32')
        return (x * mask) / (1 - self.p)


x = nd.ones((2, 5, 16, 16))
net = SpatialDropout2D(0.5)
out = net(x)

A HybridBlock is a bit trickier to implement. There can be multiple implementations, here is one of them which uses sum() to reduce the input to (batch, channel, 1, 1) shape so that a random mask can be created for each channel:

class SpatialDropout2D(gluon.HybridBlock):
    def __init__(self, p):
        super(SpatialDropout2D, self).__init__()
        self.p = p

    def hybrid_forward(self, F, x):
        if not autograd.is_training():
            return x
        mask = F.floor(F.random.uniform_like(x.sum((2, 3), keepdims=True),
                                             low=1 - self.p, high=2 - self.p))
        return F.broadcast_mul(x, mask) / (1 - self.p)


x = nd.ones((2, 5, 16, 16))
net = SpatialDropout2D(0.1)
net.hybridize()
out = net(x)

Of course you can pass batch-size and num-channels directly to the block to avoid this trick, but that’s not really convenient.

3 Likes