TensorFlow和Numpy矩阵操作中axis理解及axis=-1的解释

1. axis的基本使用

axis常常用在numpy和tensorflow中用到，作为对矩阵（张量）进行操作时需要指定的重要参数之一。设定axis=-1,0,1...，用来指定执行操作的数据如何划分。

一句话解释：设axis=i，则沿着第i个下标变化的方向进行操作！[1]

简单例子就不举了，其他博客有很多，这里举一个稍微复杂一点的三维矩阵的例子：

设embeddings是一个shape=[3,4,5]的矩阵，如下：

embeddings =    [[[-0.30166972  0.25741747 -0.07442257  0.24321035 -0.3538919 ]
                 [-0.22572032  0.1288028  -0.4686908  -0.07217035  0.05287632]
                 [ 0.15845934  0.07064888  0.00922218  0.2841002  -0.24992025]
                 [ 0.43347922 -0.43738696 -0.08176881  0.34185413 -0.2826353 ]]

[[-0.08590135  0.06792518 -0.07807922 -0.28746927 -0.10613027]
                 [ 0.07476929  0.132256   -0.0926154   0.39621904  0.2497718 ]
                 [-0.15389556  0.0867373   0.19403657 -0.11003655  0.317669  ]
                 [ 0.3949038  -0.17275128  0.34710506 -0.02576578 -0.17427891]]

[[-0.27703786  0.02631402  0.22129896 -0.07714707  0.41439041]
                 [-0.08512023  0.19059369 -0.13418713 -0.12881753 -0.26143318]
                 [-0.333749    0.27034065  0.45429572 -0.46164128 -0.3955955 ]
                 [ 0.24430516 -0.3841647   0.37126407 -0.463441   -0.1441828 ]]]

对embeddings矩阵执行下面操作：

a = tf.math.argmax(embeddings, axis=-1)   # tf.math.argmax=tf.argmax，用来返回最大数值对应的index
b = tf.math.argmax(embeddings, axis=1)
c = tf.math.argmax(embeddings, axis=0)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
   
    print(embeddings.eval())   # a.eval 在打印时，等同于 sess.run(a)
    print(a.eval())
    print(b.eval())
    print(c.eval())

得到的结果是：

[[1 1 3 0]
 [1 3 4 0]
 [4 1 2 2]]   # axis=-1, shape=[3,4]
 
[[3 0 2 3 1]
 [3 1 3 1 2]
 [3 2 2 0 0]]   # axis=1, shape=[3,5]
 
[[1 0 2 0 2]
 [1 2 1 1 1]
 [0 2 2 0 1]
 [0 1 2 0 2]]   # axis=0, shape=[4,5]

看懂了吗？参考上面的一句话解释，再结合矩阵的下标表示理解一下。

刚刚的矩阵写成下标表示就是：

embeddings = [[[a000,a001,a002,a003,a004],
              [a010,a011,a012,a013,a014],
              [a020,a021,a022,a023,a024],              
              [a030,a031,a032,a033,a034]]

[[a100,a101,a102,a103,a104],
              [a110,a111,a112,a113,a114],
              [a120,a121,a122,a123,a124],              
              [a130,a131,a132,a133,a134]]

[[a200,a201,a202,a203,a204],
              [a210,a211,a212,a213,a214],
              [a220,a221,a222,a223,a224],              
              [a230,a231,a232,a233,a234]]                

以axis=0为例,则沿着第0个下标（最左边的下标）变化的方向进行操作，也就是将除了第0个下标外，其他两个下标都相同的部分分成一组，然后再进行操作。具体分组如下（省略了一些组）：

从上图可以看出，每3个数分成一组，所以现在总共是分了4*5个组（所以最终返回的结果也是一个shape=[4,5]的矩阵），对每个组都执行一次 reduce_max操作，将每个组的三个数中数值最大的数的index返回构成矩阵即可。

这里需要特别说明一下axis=-1的操作，可能对python不熟悉的人会不理解这里的-1是哪个维度。在pyhton中，-1代表倒数第一个，也就是说，假如你的矩阵shape=[3,4,5],那么对这个矩阵来说，axis=-1，其实也就等于axis=2。因为这是个三维矩阵，所以axis可能的取值为0,1,2，所以最后一个就是2。你可以自己试试看两个取值结果是否相同。

2. 对axis的理解

通过上面的例子，你可能已经发现了，axis是将矩阵进行分组，然后再操作。而分组则意味着会降维。

以刚刚的例子，原始矩阵的shape=[3,4,5]，取axis=0再进行操作后，得到的矩阵shape=[4,5]。同样的，取axis=1再进行操作后，得到的矩阵shape=[3,5]。取axis=-1（axis=2）再操作后，shape=[3,4]。掌握这一点，能有利于你在神经网络中的变换或是数据操作中明确矩阵变换前后的形状，从而加快对模型的理解。

参考资料

[1] 这个一句话解释来源于：https://www.jb51.net/article/242077.htm

来源：https://blog.csdn.net/weixin_37821353/article/details/88367211