SoftMax

将 SoftMax 函数应用于输入张量的一个维度，生成输出张量，使得输出中的值在 \([0,1]\) 范围内，并且沿应用维度的每个切片的所有值之和等于 1。

SoftMax 函数定义如下：

\[\large{Softmax(x_{a_0,...,a_i,...,a_n}) = \frac{e^{x_{a_0,...,a_i,...,a_n}}}{\sum\limits_{j \in [:,...,a_i,...,:]}{e^{x_j}}}}\]

属性

axes 计算 SoftMax 的轴，表示为位掩码。如果输入具有 n 维度，则 \(axes < n\)。

输入

input：类型为 T 的张量

输出

output：类型为 T 的张量

数据类型

T：float16、float32、bfloat16

形状信息

input 和 output 是形状为 \([a_0,...,a_n]\) 的张量。

DLA 支持

支持 DLA FP16。

示例

SoftMax

in1 = network.add_input("input1", dtype=trt.float32, shape=(1, 1, 2, 2))
layer = network.add_softmax(in1)
layer.axes = 1 << 2  # run softmax along h dim
network.mark_output(layer.get_output(0))

inputs[in1.name] = np.array(
    [
        [
            [1.0, 1.0],
            [3.0, 4.0],
        ]
    ]
)

outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array(
    [
        [
            [0.119, 0.048],
            [0.881, 0.952],
        ]
    ],
)

C++ API

有关 C++ ISoftMaxLayer 算子的更多信息，请参阅 C++ ISoftMaxLayer 文档。

Python API

有关 Python ISoftMaxLayer 算子的更多信息，请参阅 Python ISoftMaxLayer 文档。