TensorFlow.js tf.softmax()函数

在 TensorFlow.js 中，tf.softmax() 函数可以用来对张量进行 softmax 操作。softmax 是一种常用的数学函数，用于将一个 K 维向量 $\mathbf{z}$ 通过 $\mathbf{p}$ 的一个 K 维向量来表示，每个元素 $\large p_i$ 都是 $z_i$ 的指数函数除以所有元素的指数函数之和。

$$ p_i = \frac{e^{z_i}}{\sum_{j=1}^{K}{e^{z_j}}}. $$

通过 softmax 函数，可以将一个任意的 $K$ 维向量 $\mathbf{z}$ 转换为一个概率分布向量 $\mathbf{p}$，满足 $\sum_{i=1}^{K}{p_i} = 1$。

语法

tf.softmax(logits, dim)

参数

• logits：要进行 softmax 操作的张量。
• dim（可选）：指定沿着哪个轴进行 softmax 操作。默认值为 -1，表示最后一个轴。

返回值

进行 softmax 操作之后的结果张量。

示例

以下代码演示了如何使用 tf.softmax() 函数对一个张量进行 softmax 操作：

const tf = require('@tensorflow/tfjs');

const logits = tf.tensor2d([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
const softmaxResult = tf.softmax(logits);

softmaxResult.print();

输出：

Tensor
    [[0.09003, 0.24473, 0.66524],
     [0.09003, 0.24473, 0.66524]]

在上面的示例中，我们使用了一个 $2 \times 3$ 的张量来演示 softmax 的操作。经过 softmax 操作之后，原张量被转换为了形状相同的概率分布张量。

备注

需要注意的是，由于 softmax 操作依赖于指数函数的计算，因此当输入的张量中存在非常大或非常小的值时，softmax 操作可能导致数值精度问题。在实际应用中，可以通过对输入进行缩放或对输出进行修剪等方式来规避此类问题。