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:
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
, bfloat16
Shape Information¶
input and output are tensors with a shape of \([a_0,...,a_n]\).
DLA Support¶
DLA FP16 is supported.
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.