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.