深入浅出Pytorch函数——torch.softmax/torch.nn.functional.softmax

分类目录：《深入浅出Pytorch函数》总目录
相关文章：
· 机器学习中的数学——激活函数：Softmax函数
· 深入浅出Pytorch函数——torch.softmax/torch.nn.functional.softmax
· 深入浅出Pytorch函数——torch.nn.Softmax

---

将Softmax函数应用于沿dim的所有切片，并将重新缩放它们，使元素位于$[0,1]$的范围内并和为1。

语法

torch.softmax(input, dim, *, dtype=None) -> Tensor
torch.nn.functional.softmax(input, dim=None, _stacklevel=3, dtype=None) -> Tensor

参数

• input：[Tensor] 输入张量
• dim：[int] Softmax函数将沿着dim轴计算，即沿dim的每个切片的和为1
• dtype：[可选, torch.dtype] 想要返回张量的数据类型。如果指定，则在执行操作之前将输入张量强制转换为dtype。这对于防止数据类型溢出非常有用。默认值为None

返回值

与input具有相同形状且值在[0，1]范围内的Tensor。

实例

>>> x = torch.randn(4, 5)
>>> torch.nn.functional.softmax(x, dim=0)
tensor([[0.4259, 0.5448, 0.1935, 0.3904, 0.1963],[0.1370, 0.1053, 0.1966, 0.2625, 0.4343],[0.0540, 0.2823, 0.5101, 0.2082, 0.0905],[0.3832, 0.0676, 0.0998, 0.1390, 0.2789]])>>> torch.nn.functional.softmax(x, dim=1)
tensor([[0.0728, 0.1720, 0.1233, 0.5541, 0.0779],[0.1139, 0.0667, 0.2677, 0.4361, 0.1156],[0.2576, 0.0811, 0.0671, 0.2975, 0.2968],[0.1874, 0.0358, 0.2240, 0.4470, 0.1057]])

函数实现

def softmax(input: Tensor, dim: Optional[int] = None, _stacklevel: int = 3, dtype: Optional[DType] = None) -> Tensor:r"""Applies a softmax function.Softmax is defined as::math:`\text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}`It is applied to all slices along dim, and will re-scale them so that the elementslie in the range `[0, 1]` and sum to 1.See :class:`~torch.nn.Softmax` for more details.Args:input (Tensor): inputdim (int): A dimension along which softmax will be computed.dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor.If specified, the input tensor is casted to :attr:`dtype` before the operationis performed. This is useful for preventing data type overflows. Default: None... note::This function doesn't work directly with NLLLoss,which expects the Log to be computed between the Softmax and itself.Use log_softmax instead (it's faster and has better numerical properties)."""if has_torch_function_unary(input):return handle_torch_function(softmax, (input,), input, dim=dim, _stacklevel=_stacklevel, dtype=dtype)if dim is None:dim = _get_softmax_dim("softmax", input.dim(), _stacklevel)if dtype is None:ret = input.softmax(dim)else:ret = input.softmax(dim, dtype=dtype)return ret