从LayerNorm中看keepdim参数的作用

Handsome_zhou

Transformer的LayerNorm层中做均值计算等运算时用到了keepdim参数：

1. class LayerNorm(nn.Module):
   
2.     "Construct a layernorm module (see citation for details)."
   
3.     def __init__(self, features, eps=1e-6):
   
4.         super(LayerNorm, self).__init__()
   
5.         self.a_2 = nn.Parameter(torch.ones(features))
   
6.         self.b_2 = nn.parameter(torch.zeros(features))
   
7.         self.eps = eps
   
8.    
   
9.     def forward(self, x):
   
10.         mean = x.mean(-1, keepdim=True)
    
11.         std = x.std(-1, keepdim=True)
    
12.         return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

复制代码

keepdim参数的作用是可以使得在做归并操作时原来的数据维度不变。

示例：

1. >>>a = torch.arange(12).reshape(3, 4)
   
2. >>>print(a)
   
3. tensor([[ 0,  1,  2,  3],
   
4.         [ 4,  5,  6,  7],
   
5.         [ 8,  9, 10, 11]])
   
6. 
   
7. >>>print(torch.sum(a, dim=0, keepdim=True))
   
8. >>>print(torch.sum(a, dim=0, keepdim=True).shape)
   
9. tensor([[12, 15, 18, 21]])
   
10. torch.Size([1, 4])
    
11. 
    
12. >>>print(torch.sum(a, dim=0, keepdim=False))
    
13. >>>print(torch.sum(a, dim=0, keepdim=False).shape)
    
14. tensor([12, 15, 18, 21])
    
15. torch.Size([4])

复制代码

使用 keepdim=True后，输出张量的维度不变。
keepdim = False后，输出张量丢失了第一个维度。