tf.nn.softmax(x, axis)里axis起什么作用？

在处理多分类问题时，tf.nn.softmax(x, axis) 函数是一定要使用的，那么这里的 axis 到底有什么用呢？

二维数组

首先，我们来看在二维数组中的情况：

A = np.array([[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]])
A = tf.cast(A, tf.float32)
A

<tf.Tensor: id=135, shape=(4, 3), dtype=float32, numpy=
array([[1., 2., 3.],
       [1., 2., 3.],
       [1., 2., 3.],
       [1., 2., 3.]], dtype=float32)>

当我们将这个数组输入 softmax 层后，我们改变 axis 来查看输出结果的不同：

tf.nn.softmax(A, axis=0)

<tf.Tensor: id=360, shape=(4, 3), dtype=float32, numpy=
array([[0.25, 0.25, 0.25],
       [0.25, 0.25, 0.25],
       [0.25, 0.25, 0.25],
       [0.25, 0.25, 0.25]], dtype=float32)>

tf.nn.softmax(B, axis=1)

<tf.Tensor: id=361, shape=(4, 3), dtype=float32, numpy=
array([[0.09003057, 0.24472848, 0.66524094],
       [0.09003057, 0.24472848, 0.66524094],
       [0.09003057, 0.24472848, 0.66524094],
       [0.09003057, 0.24472848, 0.66524094]], dtype=float32)>

由此可见，当在第一个维度上使用 softmax 层时，输出数组的计算过程如下：
$$\frac{e}{e+e+e+e}=0.25$$
$$\frac{e^2}{e^2+e^2+e^2+e^2}=0.25$$
$$\frac{e^3}{e^3+e^3+e^3+e^3}=0.25$$
当在第二个维度上使用 softmax 层时，输出数组的计算过程如下：
$$\frac{e}{e+e^2+e^3}=0.09003057$$
$$\frac{e^2}{e+e^2+e^3}=0.24472848$$
$$\frac{e^3}{e+e^2+e^3}=0.66524094$$

三维数组

然后，我们创建一个三维数组来进行测试：

A = np.array([[[1, 2, 3], [1, 2, 3], [1, 2, 3]], [[2, 3, 5], [2, 3, 5], [2, 3, 5]]])
A = tf.cast(A, tf.float32)
A

<tf.Tensor: id=266, shape=(2, 3, 3), dtype=float32, numpy=
array([[[1., 2., 3.],
        [1., 2., 3.],
        [1., 2., 3.]],

       [[2., 3., 5.],
        [2., 3., 5.],
        [2., 3., 5.]]], dtype=float32)>

令 axis=0：

tf.nn.softmax(A, axis=0)

<tf.Tensor: id=392, shape=(2, 3, 3), dtype=float32, numpy=
array([[[0.26894143, 0.26894143, 0.11920291],
        [0.26894143, 0.26894143, 0.11920291],
        [0.26894143, 0.26894143, 0.11920291]],

       [[0.7310586 , 0.7310586 , 0.880797  ],
        [0.7310586 , 0.7310586 , 0.880797  ],
        [0.7310586 , 0.7310586 , 0.880797  ]]], dtype=float32)>

此时我们可以这样理解，将数组写成：

[[[1., 2., 3.], [1., 2., 3.], [1., 2., 3.]],
 [[2., 3., 5.], [2., 3., 5.], [2., 3., 5.]]]

然后按照第一维度计算：
$$\frac{e}{e+e^2}=0.26894143$$
$$\frac{e^2}{e^2+e^3}=0.26894143$$
$$\dots \dots$$
令 axis=1：

tf.nn.softmax(A, axis=1)

<tf.Tensor: id=423, shape=(2, 3, 3), dtype=float32, numpy=
array([[[0.33333334, 0.33333334, 0.33333334],
        [0.33333334, 0.33333334, 0.33333334],
        [0.33333334, 0.33333334, 0.33333334]],

       [[0.33333334, 0.33333334, 0.33333334],
        [0.33333334, 0.33333334, 0.33333334],
        [0.33333334, 0.33333334, 0.33333334]]], dtype=float32)>

此时的计算规则是：将整个数组分成两个 $3\times 3$ 的小数组，这两个小数组在之后的运算中毫无关联。这与 axis=0 时的运算规则有很大的区别，因为当 axis=0 时，所有的计算都有这两个小数组共同参与。
$$\frac{e}{e+e+e}=0.33333334$$
$$\frac{e^2}{e^2+e^2+e^2}=0.33333334$$
$$\dots \dots$$
$$\frac{e^5}{e^5+e^5+e^5}=0.33333334$$
令 axis=2：

tf.nn.softmax(A, axis=2)

<tf.Tensor: id=424, shape=(2, 3, 3), dtype=float32, numpy=
array([[[0.09003057, 0.24472848, 0.66524094],
        [0.09003057, 0.24472848, 0.66524094],
        [0.09003057, 0.24472848, 0.66524094]],

       [[0.04201007, 0.11419519, 0.8437947 ],
        [0.04201007, 0.11419519, 0.8437947 ],
        [0.04201007, 0.11419519, 0.8437947 ]]], dtype=float32)>

此时的计算规则与 axis=1 时的运算规则相似，两个小数组之间仍然没有关系。
$$\frac{e}{e+e^2+e^3}=0.09003057$$
$$\frac{e^2}{e+e^2+e^3}=0.24472848$$
$$\frac{e^3}{e+e^2+e^3}=0.66524094$$
$$\frac{e^2}{e^2+e^3+e^5}=0.04201007$$
$$\frac{e^3}{e^2+e^3+e^5}=0.11419519$$
$$\frac{e^5}{e^2+e^3+e^5}=0.8437947$$