coordConv layer

Thank you @kobenaxie, great implementation, I haven’t noticed the channel layers normalization until I saw your solution.

@safrooze thank you, I’ve learned so much from your implementation. Especially the trick to override forward to get the shape of the array is super important. It will prove helpful in so many future applications (you should see the hack I had to perform, with successive splits, to implement in a HybridBlock format the pyramid scene pooling operator). Initially I tried something like you proposed, and used F.shape_array to get the shape of the array and use F.arange to create the range of the indices. However, I couldn’t extract the actual values of the F.shape_array unless I used asscalar() or asnumpy() - which I wanted to avoid for performance reasons. For example when I was trying shape = F.shape_array(x) then (e.g.) shape[0] was an NDArray instance (not an int), and I couldn’t use it in F.arange(start = shape[0], ...).

Again, many thanks!

edit minor modification to your code (the return statement in the forward override wasn’t working for me)

class coordConv2D(HybridBlock):
    def __init__(self, channels, kernel_size=3, padding=1, strides=1, **kwards):
        HybridBlock.__init__(self, **kwards)
        self.xs = None  # x.shape, initialized in forward()
        with self.name_scope():
            self.conv = gluon.nn.Conv2D(channels=channels, kernel_size=kernel_size, padding=padding, strides=1,
                                        **kwards)

            
    def _ctx_kwarg(self,x):
        if isinstance(x, nd.NDArray):
            return {"ctx": x.context}
        return {}

    def _coord_array(self,F, start, stop, rows, cols, batch_size, **ctx_kwarg):
        """ Outputs (rows, cols) with each NDArray in NCHW layout (C is 1) """
        row_array = F.arange(start=start, stop=stop, step=(stop - start) / rows, **ctx_kwarg)
        col_array = F.arange(start=start, stop=stop, step=(stop - start) / cols, **ctx_kwarg)
        coord_rows = F.repeat(F.reshape(row_array, (1, 1, rows, 1)), repeats=cols, axis=3)
        coord_cols = F.repeat(F.reshape(col_array, (1, 1, 1, cols)), repeats=rows, axis=2)
        return (F.repeat(coord_rows, repeats=batch_size, axis=0),
                F.repeat(coord_cols, repeats=batch_size, axis=0))
        
    def forward(self, x):
        """ Override forward to collect shape in hybridized mode """
        # x is in NCHW layout
        self.xs = x.shape
        #return super(CoordConv, self).forward(x) # This wasn't working on my machine - python3.6 
        return HybridBlock.forward(self,x)

    def hybrid_forward(self, F, x):
        rows, cols = self._coord_array(F, -1, 1, rows=self.xs[2], cols=self.xs[3],
                                  batch_size=self.xs[0], **self._ctx_kwarg(x))
        x = F.concat(x, rows, cols, dim=1)
        return self.conv(x), rows, cols  # I am also returning rows, cols for vis purposes