Defining a layer in the __init__ method of the model

1

This is a more general question but when is it ok to define just a layer and use it multiple times ? For example, if I want to use batchnorm in my model, can I just define

self.batchnorm = nn.BatchNorm2d(num_features, eps=1e-05, momentum=0.1, affine=True)

and then just apply self.batchnorm after every layer? Or should I define a separate batchnorm layer for every time I use it, for example

self.batchnorm1 = nn.BatchNorm2d(num_features, eps=1e-05, momentum=0.1, affine=True)
self.batchnorm2 = nn.BatchNorm2d(num_features, eps=1e-05, momentum=0.1, affine=True)
self.batchnorm3 = nn.BatchNorm2d(num_features, eps=1e-05, momentum=0.1, affine=True)
2

I think it would in general depend on if the module is stateless or not; In the BatchNorm case you’d be sharing weights, bias, running_mean, running_var, which is probably not what you want.