详解TensorFlow的 tf.nn.softmax 函数：softmax 激活函数

什么是softmax函数

在机器学习的过程中，很多训练算法都是基于概率论的基础理论进行的。softmax函数是一种用于归一化多维向量的函数，通常作为神经网络的输出层的激活函数，可以将任意实数值向量转换为概率分布。

softmax函数的数学定义

对于给定的一个包含m个元素的向量 $\boldsymbol{z}=(z_1,z_2,...,z_m)$，softmax函数通过以下公式定义：

$$
\text{softmax}(\boldsymbol{z})_j = \frac{e^{zj}}{\sum\limits{i=1}^{m} e^{z_i}}, \ \ 1 \leq j \leq m,
$$

其中，分母用于保证所有计算后的分量和为1，分子 $e^{z_j}$ 是将原始分量扩大的一个操作，其实就是将分量转化成概率分布。

TensorFlow中 tf.nn.softmax 函数使用方法

TensorFlow中的tf.nn.softmax函数可以直接对张量进行softmax操作，函数的完整定义如下：

tf.nn.softmax(logits, axis=None, name=None)

函数的参数说明如下：

• logits：需要进行softmax操作的张量，可以任意维度。
• axis：指定进行softmax操作的维度，默认为None，此时将进行最后一维的softmax操作。

示例1：对一个简单的二维张量进行softmax操作

import tensorflow as tf

x = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
y = tf.nn.softmax(x, axis=1)
print(y)

输出结果：

tf.Tensor(
[[0.09003057 0.24472848 0.66524094]
 [0.09003057 0.24472848 0.66524094]], shape=(2, 3), dtype=float32)

示例2：对MNIST手写数字数据集的预测结果进行softmax操作

import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

# 载入MNIST数据集
mnist = input_data.read_data_sets('MNIST_data', one_hot=True)

# 构建模型
x = tf.placeholder(tf.float32, [None, 784])
W = tf.Variable(tf.zeros([784, 10]))
b = tf.Variable(tf.zeros([10]))
y = tf.nn.softmax(tf.matmul(x, W) + b)

# 运行模型
with tf.Session() as sess:
    sess.run(init)
    # 获取测试集的预测结果
    correct_prediction = tf.equal(tf.argmax(y, 1), tf.argmax(y_, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    print("Accuracy:", sess.run(accuracy, feed_dict={x: mnist.test.images, y_: mnist.test.labels}))

总结

softmax函数是机器学习中常用的归一化函数，可以将任何一个m维向量转换成概率分布。TensorFlow提供了tf.nn.softmax函数，可以在计算图中直接进行softmax操作。在构建深度学习模型时，经常需要使用softmax函数作为输出层的激活函数，将卷积神经网络的最后一层输出的向量转换成分类概率分布。