I know LayerNorm and InstanceNorm1d can do normalization over the last dimension. If I want to do it over the first dimension, I have to transpose the input tensor or calculate the mean and standard-deviation by my self, which would consume much time and memory.
if I do it by my self or transpose the input:
class Net1(nn.Module):
def __init__(self):
super(Net1, self).__init__()
def forward(self, x):
return (x - x.mean(dim=0, keepdim=True))/ (x.var(dim=0, unbiased=False, keepdim=True) + 1e-5).sqrt()
class Net2(nn.Module):
def __init__(self):
super(Net2, self).__init__()
def forward(self, x):
return nn.functional.layer_norm(x.transpose(0, 2), (1000,)).transpose(0, 2)
class Net3(nn.Module):
def __init__(self):
super(Net3, self).__init__()
def forward(self, x):
return nn.functional.instance_norm(x.transpose(0, 2)).transpose(0, 2)
net1 = Net1()
net2 = Net2()
net3 = Net3()
x = torch.randn(1000, 100000, 10, requires_grad=True)
%timeit net1(x) #4.7 s ± 32.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit net2(x) #2.71 s ± 46.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit net3(x) #3.33 s ± 82.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit net1(x).mean().backward() #22.1 s ± 52.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit net2(x).mean().backward() #11.3 s ± 169 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit net3(x).mean().backward() #10.3 s ± 134 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
But if I do normalization over the last dimension without transpose, I can save half of time.
y = torch.randn(10, 100000, 1000, requires_grad=True)
class Net5(nn.Module):
def __init__(self):
super(Net5, self).__init__()
def forward(self, x):
return nn.functional.layer_norm(x, (1000,))
class Net6(nn.Module):
def __init__(self):
super(Net6, self).__init__()
def forward(self, x):
return nn.functional.instance_norm(x)
net5 = Net5()
net6 = Net6()
%timeit net5(y) #1.24 s ± 22.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit net6(y) #1.73 s ± 16.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit net5(y).mean().backward() #4.07 s ± 28.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit net6(y).mean().backward() #4.69 s ± 32.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Is there any method to do it over specific dimension effectively?