Get HybridBlock layer shape on runtime

Gluon
6

It’s really simple. Let’s say you have a custom pooling layer that is not hybridizable:

class CustomPooling(Block):
    def __init__():
        # Initialize
    def forward(x):
        # custom pooling op
        return y

Now if everything before this layer are hybridizable, then put them into a single hybrid block. Same with all layers after this custom pooling layer:

class BeforePooling(HybridBlock):
    def __init__():
        # Initialize
    def hybrid_forward(F, x):
        # perform ops
        return y

class AfterPooling(HybridBlock):
    def __init__():
        # Initialize
    def hybrid_forward(F, x):
        # perform ops
        return y

Now you can either create a custom block class, or use nn.sequential to chain these three sections up:

net = nn.sequential()
with net.name_scope():
    net.add(BeforePooling)
    net.add(CustomPooling)
    net.add(AfterPooling)
net.hybridize()

What hybridize() does in this case is that it goes through all the children and hybridizes every child that is hybridizable. Keep in mind that the hybridization in this case is done one child as a time. So if you, for example, use nn.sequential (or a custom block) and have two hybridizable layers (e.g. two conv layers) added as two separate children (e.g. in two add() calls), even though they maybe back to back, the hybridization creates a separate symbolic graph for each child if the parent is a Block. Just make sure that you collect all your hybridizable layers under one HybridBlock parent before adding on the normal Block layers and then hybridize the final parent Block.

2 Likes