When working with vectorial data, I sometimes need to leave the batch x dimension format in favour of batch x samples x dimension.
The sampling dimension can originate from the latent space of importance weighted autoencoders, or just multiple measurements of the same instance during an experiment.
While linear layers easily broadcast along the additional dimensions, batch norm layers need special care here.
My work around about the missing bach norm broadcasting is a double reshape
from reshape(batch * samples x dimension) => batch_norm(...) => reshape(batch, samples, dimension
Is there a way to extend batch norm without the supposedly inefficient reshaping?
For illustration
import torch
import torch.nn as nn

lin = nn.Linear(10, 10)

b, p, d = 16, 80, 10
x = torch.randn(b, p, d)
x_ = x.reshape(b * p, d)

o = lin(x)
o_ = lin(x_)

print(torch.allclose(o, o_.reshape(b, p, d)))
">> yay :)"


bn = nn.BatchNorm1d(10)
bn(x_) # is ok
bn(x)
">> nay :?"
"RuntimeError: running_mean should contain 80 elements not 10"