oneflow.nn.Softmax¶
-
class
oneflow.nn.
Softmax
(dim: Optional[int] = None)¶ Applies the Softmax function to an n-dimensional input Tensor rescaling them so that the elements of the n-dimensional output Tensor lie in the range [0,1] and sum to 1.
Softmax is defined as:
\[\text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}\]When the input Tensor is a sparse tensor then the unspecifed values are treated as
-inf
.- Shape:
Input: \((*)\) where * means, any number of additional dimensions
Output: \((*)\), same shape as the input
- Returns
a Tensor of the same dimension and shape as the input with values in the range [0, 1]
- Parameters
dim (int) – A dimension along which Softmax will be computed (so every slice along dim will sum to 1).
For example:
>>> import numpy as np >>> import oneflow as flow >>> m = flow.nn.Softmax(dim = 2) >>> x = flow.Tensor( ... np.array( ... [[[-0.46716809, 0.40112534, 0.61984003], ... [-1.31244969, -0.42528763, 1.47953856]]] ... ) ... ) >>> out = m(x) >>> out tensor([[[0.1575, 0.3754, 0.4671], [0.0507, 0.1230, 0.8263]]], dtype=oneflow.float32)