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$$