SoftMax

Applies the SoftMax function on one of the dimensions of an input tensor into an output tensor, so that the values in the output lies in the range \([0,1]\), and the sum of all the values of each slice along the applied dimension is equal to 1.

The SoftMax function is defined as:

\[\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}}}}\]

Attributes

axes The axis along the SoftMax is computed, represented as a bitmask. If input has n dimensions, \(axes < n\).

Inputs

input: tensor of type T

Outputs

output: tensor of type T

Data Types

T: float16, float32

Shape Information

input and output are tensors with a shape of \([a_0,...,a_n]\).

Examples

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

For more information about the C++ ISoftMaxLayer operator, refer to the C++ ISoftMaxLayer documentation.

Python API

For more information about the Python ISoftMaxLayer operator, refer to the Python ISoftMaxLayer documentation.