How to share a single parameter across two gluon layers?

I want to use two gluon.nn.BatchNorm layers with shared gamma parameter. I know how to make the two layers share all the parameters but I do not know how to share a single parameter. Following is the code to share all the parameters.

BN1 = gluon.nn.BatchNorm()
BN2 = gluon.nn.BatchNorm(params=BN1.params)

Any help would be highly appreciated. I want these two layers to share only “gamma” parameter instead of all the parameters.


the only way I can think to do this would by to create your own custom HybridBlock for BatchNorm. This Block would depend on two other HybridBlocks that simply acts as containers for the gamma and beta parameters respectively. That way you’d be able to share parameters by doing something like:

beta1 = CustomBatchNormBeta()
gamma1 = CustomBatchNormGamma()
BN1 = CustomBatchNorm(beta1, gamma1)

beta2 = CustomBatchNormBeta()
gamma2 = CustomBatchNormGamma(params=gamma1.params)
BN2 = CustomBatchNorm(beta2, gamma2)