深度学习中的轴/axis/dim全解

心已赠人 2022-12-06 12:39 235阅读 0赞

在深度学习中，轴，指的就是张量的层级，一般通过参数axis/dim来设定。很多张量的运算、神经网络的构建，都会涉及到轴，但到底取哪个轴，却不是那么容易把握。

下面会针对轴/axis/dim，基于 Pytorch 的代码和实例，尝试去理清张量运算中轴/axis/dim的设定。

#### 轴的概念 ####

对于一个张量，它的shape有几维，就对应有几个轴，也就对应着张量的层级，最直观的可以通过看最前面的方括号数量来判断。

import torch
    a = torch.Tensor([[1,2,3], [4,5,6]])
    b = torch.Tensor([[7,8,9], [10,11,12]])
    c = torch.Tensor([[[1,2,3], [4,5,6]], [[7,8,9], [10,11,12]]])
    print(a.shape)
    
    # torch.Size([2, 3])

上面的张量 a 和 b，都对应两个轴。axis/dim=0 对应 shape \[2, 3\] 中的2，axis/dim=1 对应 shape \[2, 3\] 中的3，而张量 c 有三个轴。

张量运算时对轴参数的设定非常常见，在 Numpy 中一般是参数axis，在 Pytorch 中一般是参数dim，但它们含义是一样的。

#### 轴的使用 ####

在做张量的拼接操作时，**axis/dim设定了哪个轴，那对应的轴在拼接之后张量数会发生变化**

>> torch.cat((a,b), dim=0)
    tensor([[ 1.,  2.,  3.],
            [ 4.,  5.,  6.],
            [ 7.,  8.,  9.],
            [10., 11., 12.]])
    
    >> torch.cat((a,b), dim=1)
    tensor([[ 1.,  2.,  3.,  7.,  8.,  9.],
            [ 4.,  5.,  6., 10., 11., 12.]])

对于上面torch中的cat操作，当设定dim=0时，两个维度是(2,3)的张量合并成了一个(4,3)的张量，在第0维，张量数从2变成了4，第1维没有变化；当设定dim=1时，在第1维，张量数从3变成了6，第0维没有变化。

在做张量的运算操作时，**axis/dim设定了哪个轴，就会遍历这个轴去做运算，其他轴顺序不变**

>> torch.softmax(a, dim=0)
    tensor([[0.0474, 0.0474, 0.0474],
            [0.9526, 0.9526, 0.9526]])
    
    >> torch.softmax(a, dim=1)
    tensor([[0.0900, 0.2447, 0.6652],
            [0.0900, 0.2447, 0.6652]])

对于上面torch中的 softmax 操作，当设定 dim=0 时，就是其他轴不变，单次遍历 dim=0 轴的所有元素去做运算，上例中就相当于分别取了张量a中的第0列、第1列、第2列去做计算。

换一个角度，假设用for循环去遍历一个张量，那运算中设定的dim就是被放在最内层的for循环，其它的轴保持正常的顺序。

可以用下面的例子作为验证，这里tensor c 的shape 是 (m,n,p)，用for循环去计算 torch.softmax(c, dim=1)

# for循环计算方式
    c = torch.Tensor([[[1,2,3], [4,5,6]], [[7,8,9], [10,11,12]]])	# shape (2,2,3)
    m,n,p = c.shape
    res = torch.zeros((m,n,p))
    for i in range(m):
        for j in range(p):
            res[i,:,j] = torch.softmax(torch.tensor([c[i,k,j] for k in range(n)]), dim=0)  #这里对应最内层的for循环
    
    # 库函数设定轴计算方式
    res1 = torch.softmax(c, dim=1)
    print(res.equal(res1))		# True

axis/dim使用小总结：

1.  在做张量的拼接操作时，**axis/dim设定了哪个轴，那对应的轴在拼接之后张量数会发生变化**
2.  在做张量的运算操作时，**axis/dim设定了哪个轴，就会遍历这个轴去做运算，其他轴顺序不变**

实际上，第一条拼接操作也可以用第二条去理解，但拼接的轴张量数会发生变化更好理解和记忆。

#### 轴的实例 ####

其实一个轴设定的变化，会带来很大的差异，最典型的就是 BatchNorm 和 LayerNorm 了。  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3Nvbmd5dW5saTExMTE_size_16_color_FFFFFF_t_70_pic_center]

BatchNorm 和 LayerNorm 是针对数据的不同轴去做norm，假设输入数据的维度是(N,C,H,W)，分别对应batch数，核数，高，宽，BatchNorm 就对应dim=0，LayerNorm 就对应dim=1，在不考虑移动平均等具体细节问题时，两者在形式上可以统一，只有一个dim参数的差别。

Pytorch 的实现（简化版）如下：

class Norm(nn.Module):
        def __init__(self, num_features, variance_epsilon=1e-12):
            super(Norm, self).__init__()
            self.gamma = nn.Parameter(torch.ones(num_features))
            self.beta = nn.Parameter(torch.zeros(num_features))
            self.variance_epsilon = variance_epsilon    # 一个很小的常数，防止除0
    
        def forward(self, x, dim):
            u = x.mean(dim, keepdim=True)
            s = (x - u).pow(2).mean(dim, keepdim=True)
            x_norm = (x - u) / torch.sqrt(s + self.variance_epsilon)
            return self.gamma * x_norm + self.beta

当然，不仅仅是在深度学习里面，在 Numpy，Pandas中，轴的设定都经常会遇到，但它们都是相通的，希望本文能帮你更好的理解它 —> 轴。

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3Nvbmd5dW5saTExMTE_size_16_color_FFFFFF_t_70_pic_center]: /images/20221123/baa5ba20f5774ccc88c51af674668654.png