Custom layer - infer shape after first forward pass

Gluon
11

Hi to all,

@safrooze gave the solution in this topic in the discussion forum. The trick is to overwrite the forward function, and getting the layer shape in there. Example

from mxnet import gluon

class GetShape(gluon.HybridBlock):
    def __init__(self,nchannels=0, kernel_size=(3,3), **kwards):
        gluon.HybridBlock.__init__(self,**kwards)
        
        self.layer_shape = None
        
        with self.name_scope():
            self.conv = gluon.nn.Conv2D(nchannels,kernel_size=kernel_size)
            
            
            
    def forward(self,x):
        self.layer_shape = x.shape
        
        return gluon.HybridBlock.forward(self,x)
    
    def hybrid_forward(self,F,x):
        print (self.layer_shape)
        out = self.conv(x)
        return out

mynet = GetShape(nchannels=12)
mynet.hybridize()

mynet.initialize(mx.init.Xavier(),ctx=ctx)
xx = nd.random.uniform(shape=[32,8,128,128])
out = mynet(xx)
# prints (32, 8, 128, 128)

Thank you @safrooze !!