SoftmaxOperation

Applicable Products

Hardware Model	Supported	Remarks
Atlas 350 accelerator cards	√	-
Atlas A3 inference products / Atlas A3 training products	√	-
Atlas A2 training products / Atlas A2 inference products	√	-
Atlas training products	√	-
Atlas inference products	√	-
Atlas 200I/500 A2 inference products	√	Only the float and float16 data types are supported.

Description

Softmax is a normalization activation function that normalizes the data of one or more dimensions of a tensor. The range of each element in the corresponding dimension is (0, 1) and the sum is 1.

Figure 1 SoftmaxOperation

Formula

$\text{[math]}$

Application Scenarios

It is usually used as the last layer of the model. For example, in the Transformer model, softmax is performed on K and Q of Self-Attention to obtain the corresponding correlation.

$\text{[math]}$

Definition

      
           struct SoftmaxParam {
    SVector<int64_t> axes;
    uint8_t rsv[8] = {0};
};

Parameters

Member	Type	Default Value	Description
axes	SVector<int64_t>	-	Specified axis (dimension). Multiple axes are supported. The value of axes cannot be empty. When multiple axes are specified, they must be consecutive and arranged in ascending order. The element of axes must be greater than or equal to -1 and less than the dimension of the input x. -1 can only exist independently. For example, [-1,0] is not supported.
rsv[8]	uint8_t	{0}	Reserved

Member

Type

Default Value

Description

axes

SVector<int64_t>

Specified axis (dimension). Multiple axes are supported.

The value of axes cannot be empty. When multiple axes are specified, they must be consecutive and arranged in ascending order.
The element of axes must be greater than or equal to -1 and less than the dimension of the input x. -1 can only exist independently. For example, [-1,0] is not supported.

rsv[8]

uint8_t

{0}

Reserved

Input

Parameter	Dimension	Data Type	Format	Description
x	[dim_0, dim_1, ..., dim_n]	float16/float/bf16	ND	Input tensor.

Output

Parameter	Dimension	Data Type	Format	Description
output	[dim_0, dim_1, ..., dim_n]	float16/float/bf16	ND	Output tensor. Has the same dimension and data type as the input tensor.

Functions

One-dimensional softmax
When the axes parameter contains only one element, the softmax operation is performed on the corresponding dimension of the input tensor. axes=[x] indicates the following equivalent operation in PyTorch:

$\text{[math]}$
Multiple-dimensional softmax
When axes contains multiple elements and meets the requirements, the softmax operation applies to multiple dimensions of the input tensor. When axes=[x, x+1], the softmax operation is performed on the elements corresponding to the x and x+1 dimensions. Tile multiple consecutive dimensions in axes into one dimension, perform the softmax operation on the dimension, and reshape the dimension back to the original shape.

Parent topic: atb/infer_op_params.h