SoftMax
场景说明
多分类任务中沿指定维度计算归一化指数概率分布，目前kdnn支持torch.float32数据类型，其他数据类型会走开源分支。
示例代码
1 2 3 4 5 6 7 8 9 10 11 12 13 | import torch import torch.nn as nn #使能KDNN torch._C._set_kdnn_enabled(True) # 构造Softmax层对列做归一化 m = nn.Softmax(dim=1) logits = torch.randn(4, 5) # 4个样本，5分类 probs = m(logits) # 每行概率和为1 print(probs) |
父主题： 使用示例