pytorch softmax函数使用

在PyTorch中，可以使用torch.nn.functional.softmax函数来实现softmax操作。该函数的输入可以是任何张量，其中最后一个维度将被视为类别维度，并将在该维度上应用softmax操作。下面是一个使用PyTorch的softmax函数的示例代码：

import torch
import torch.nn.functional as F

input_tensor = torch.randn(2, 3, 4)  # 2个样本，3个类别，每个类别有4个特征

# 沿着最后一维计算softmax
output_tensor = F.softmax(input_tensor, dim=-1)

print(output_tensor)

在这个例子中，我们创建了一个形状为(2, 3, 4)的张量作为softmax函数的输入。在张量的最后一维上应用了softmax操作，并得到了一个形状相同的输出张量。输出张量中的每个元素都是在对应输入张量元素所在类别的条件下的概率值，因此可以被看作是一个概率分布。