SoftmaxOperation
Description
Softmax is a normalization activation function that normalizes data in one or more dimensions of a tensor. The range of each element is between (0, 1), and the sum of one or more dimension elements is 1.
Formula
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.
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.
|
|
rsv[8] |
uint8_t |
{0} |
Reserved |
Input
|
Parameter |
Dimension |
Data Type |
Format |
Description |
|---|---|---|---|---|
|
x |
[-1,…,-1] The value -1 indicates that the size of the current dimension is not restricted. |
float16/float/bf16 |
ND |
Input tensor. |
Output
|
Parameter |
Dimension |
Data Type |
Format |
Description |
|---|---|---|---|---|
|
output |
[-1,…,-1] The value -1 indicates that the size of the current dimension is not restricted. |
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. When axes=[x], the softmax operation is equivalent to PyTorch's.
- Multiple-dimensional softmax
When the axes parameter contains multiple elements and meets the requirements, the softmax operation is performed on 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.
