numpy和tensorflow中的关于参数axis的辅助理解方法

发表于   |   更新于 | 分类于 | 评论数： | 阅读次数：
本文字数： 2.9k | 阅读时长 ≈ 3 分钟

一句话总结：固定除axis之外的其他维度，在axis指定的维度上进行函数操作，轴的计数按shape的形状从左往右为0轴，1轴，2轴…（详解请看下面的推导）

首先声明：axis的默认值不是0，这一点我发现很多博客文章都搞错了。所以一定要知道，axis的默认值不是0，0代表0轴，而默认值是将整个shape拉伸成一个一维向量，在整个向量上求解过

axis的默认值不是0

当给axis赋值为0时，和采取默认值时的表现是完全不同的，从下面的代码就可以看出。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
>>> z #大小为2×3×4的数组
array([[[ 2,  3,  4,  8],
        [ 3,  1,  4,  1],
        [ 6,  3,  2,  6]],

       [[10,  2, 45,  2],
        [ 2,  4,  5, 10],
        [22,  4,  4,  1]]])
>>> np.sum(z,axis=0)  # axis=0
array([[12,  5, 49, 10],
       [ 5,  5,  9, 11],
       [28,  7,  6,  7]])
>>> np.sum(z)   #axis不指定，取默认值
154

理解axis参数的作用

刚开始学习numpy和tensorflow的朋友经常遇到类似下面这样的一些函数：

1
2
3
4
5
6
7
8
9
10
11
12
#python
x=[[1,2],[5,1]]
x=np.array(x)
z1=np.max(x,axis=0)
z2=np.max(x,axis=1)


#tensorflow
x=tf.constant([[1.,2.],[5.,2.]])  
x=tf.shape(x)  
z1=tf.reduce_max(x,axis=0)#沿axis=0操作  
z2=tf.reduce_max(x,axis=1)#沿axis=1操作

类似的还有argmax，sum等等函数，它们都含有一个名为axis的参数，那这个参数到底是什么意思呢？一句话总结就是：沿着axis指定的轴进行相应的函数操作

直接看这句话可能看不懂，下面用一个最简单的例子来说明一下。

1
2
3
4
5
6
7
8
9
10
11
12
import numpy as np
#首先，创建一个2×3维的numpy的array数组
x=[[2,3,4],[1,2,5]]
x=np.array(x)
#然后，计算不同参数下np.max的输出

print(np.max(x))
# 5
print(np.max(x,0))
# [2,3,5]
print(np.max(x,1))
# [4,5]

可以看到，如果不知道axis，那么默认就是取得整个数组的最大值，这相当于把多维数组展开成一维，然后找到这个一维数组里的最大值。
而当axis=0时，直观上来看就是取得每一列的最大值，源数组总共为2行3列，所以最终的输出包含3个元素。
当axis=1时，就相当与是取每一行的最大值。

上面的理解方式在二维数组还比较直观，但是如果数组达到3维4维甚至更高维时，就不能简单的从行列角度出发去理解了，这时应该考虑从“轴”的角度来看。首先，明确一点，“轴”是从外向里的，也就是说，最外层的是0轴，往内一次是1轴，2轴… 。 具体可以看下面的例子：

1
2
3
4
5
6
7
8
9
10
>>> z
array([[[ 2,  3,  4,  8],
        [ 3,  1,  4,  1],
        [ 6,  3,  2,  6]],

       [[10,  2, 45,  2],
        [ 2,  4,  5, 10],
        [22,  4,  4,  1]]])
>>> z.shape
(2, 3, 4)

可以看到，这是一个2×3×4的三位数组，其中0轴对应第一维（2），1轴对应第二维（3），2轴对应第三维（4）。当我们指定了函数按某一轴来计算时，函数的输出数组的shape就是去掉当前轴的shape，如下所示。

1
2
3
4
5
6
>>> np.max(z,axis=0).shape
(3, 4)
>>> np.max(z,axis=1).shape
(2, 4)
>>> np.max(z,axis=2).shape
(2, 3)

而对于输出数组的每一个元素output[i][j]的值，实际上就是z[i][...][j]集合中的最大值，如下面的代码所示。其中当axis=0时，输出数组output的shape为3×4，其中output.[2][3]的值，实际上就是z[0][2][3],z[1][2][3]的最大值，也就是（6，1）中的最大值，即为output.[2][3]=6。

再如axis=1时，输出数组output的shape为2×4，其中output.[1][2]的值，实际上就是z[1][0][2],z[1][1][2],z[1][2][2]中的最大值，也就是（45，5，4）中的最大值，即为output.[1][2]=45]。

1
2
3
4
5
6
7
8
9
10
>>> np.max(z,axis=0)
array([[10,  3, 45,  8],
       [ 3,  4,  5, 10],
       [22,  4,  4,  6]])
>>> np.max(z,axis=1)
array([[ 6,  3,  4,  8],
       [22,  4, 45, 10]])
>>> np.max(z,axis=2)
array([[ 8,  4,  6],
       [45, 10, 22]])

数学公式总结

用形式化的数学语言总结上面的过程就是：
对于大小为[i,j,k]的输入数组z，假设axis=0，那么输出矩阵output的大小就为[j,k]，并且output的每一个元素的计算方式如下：

如果axis=1，那么输出矩阵output的大小就为[i,k]，并且output的每一个元素的计算方式如下：

对于4维，5维甚至无限维的情况，计算方法是一样的，你不妨自己推导一下，如果有任何问题，欢迎可以在评论中留言。

另外，对于其他的sum，argmax等等函数中的计算方法也是一样的，只需要把函数max换成对应的函数即可，如下所示：

sum：

argmax: