SoftMax

Scenario Description

Calculate the normalized exponential probability distribution along the specified dimension in a multiclass classification task. Currently, KuDNN supports the torch.float32 data type. For other data types, see the open-source branch.

Sample Code

import torch
import torch.nn as nn

# Enable KuDNN.
torch._C._set_kdnn_enabled(True)

# Construct the SoftMax layer to normalize columns.
m = nn.Softmax(dim=1)
logits = torch.randn(4, 5) # 4 samples, 5 classes
probs = m(logits)

# The sum of probabilities in each row is 1.
print(probs)

Parent topic: Examples