SoftMax

接口功能

沿指定维度计算归一化指数概率分布。

函数原型

torch.nn.Softmax(
    dim: Optional[int] = None
)
# 或函数式
torch.nn.functional.softmax(
    input: Tensor, 
    dim: int, 
    dtype: Optional[torch.dtype] = None
) -> Tensor

参数说明

表1 参数说明
参数	类型	必填	说明
dim	int	是	计算维度（如dim=1表示对每行计算）。
input	int	是	输入张量（函数式专用）。
dtype	torch.dtype	否	输出数据类型（函数式专用）。

父主题： 算子接口