在遇到形如tf.argmax(logits, axis=-1)
的代码时,axis参数的含义非常容易令人疑惑。在二维情形下,axis=0
表示求每列的最大值的下标,axis=1
表示求每行最大值的下标。但是在更高维度下呢?
我们不妨假设数组A满足A.shape=(2,4,8,16)
,研究A生成的数组(Ax=argmax(A,axis=x)
)的shape,结果如下表所示:
A0.shape | A1.shape | A2.shape | A3.shape |
---|---|---|---|
(4,8,16) | (2,8,16) | (2,4,16) | (2,4,8) |
不难发现,Ax.shape比原来的A.shape恰好少掉了第x个值。事实上,axis=x就表示将原数组的第x维压缩,最终得到的第x维上各值的argmax。
特别地,我们知道在python中对一个长度的n的数组A,我们有A[n-i]=a[-i],同理,我们常见的axis=-1事实上就是表示压缩最后一个维度,在常见的二维数组中,就表示对压缩列,即对每行求和。
在其他一些降维的操作中,axis有着类似的作用。