PyTorch 中的 dim 参数softmax

在本文中，我们将介绍在 PyTorch 中使用 softmax 函数时的 dim 参数。softmax 是一种常用于多分类问题的激活函数，它将输入向量转换为概率分布。在 PyTorch 中，可以使用 torch.softmax() 函数来实现 softmax 操作，并且该函数提供了 dim 参数用于指定在哪个维度上进行 softmax 计算。

阅读更多：Pytorch 教程

softmax 函数简介

softmax 函数是一种常用的概率分布函数，可以将一个 n 维向量转换为具有 n 个元素的概率分布。它的定义如下：

softmax(x[i]) = exp(x[i]) / sum(exp(x[i]))

其中 x 为输入向量，i 表示向量的第 i 个元素。softmax 函数的输出范围在 0 到 1 之间，并且所有输出值的和为 1。

PyTorch 中的 softmax 函数

在 PyTorch 中，可以使用 torch.softmax() 函数来对输入进行 softmax 操作。该函数的调用方式如下：

torch.softmax(input, dim=None, dtype=None)

其中 input 是输入的张量，dim 是进行 softmax 计算的维度。如果不指定 dim 参数，则默认对最后一个维度进行 softmax 计算。

示例说明

示例 1

假设我们有一个大小为 (3, 4) 的张量 tensor，表示一个三个样本四个类别的预测结果。我们希望对每个样本的预测结果进行 softmax 操作。

import torch

tensor = torch.tensor([[0.1, 0.2, 0.3, 0.4],
                       [0.2, 0.3, 0.4, 0.1],
                       [0.3, 0.1, 0.2, 0.4]])

softmax_tensor = torch.softmax(tensor, dim=1)

print(softmax_tensor)

运行以上代码，得到的输出为：

tensor([[0.2138, 0.2341, 0.2709, 0.2813],
        [0.2498, 0.2733, 0.2999, 0.1771],
        [0.2627, 0.1864, 0.2169, 0.3339]])

上述代码通过指定 dim=1，对每个样本的预测结果进行 softmax 计算，得到了每个样本的类别概率分布。

示例 2

现在，假设我们有一个大小为 (2, 3, 4) 的张量 tensor，表示两个样本的三个特征在四个类别上的得分。我们希望对每个特征在每个样本的类别得分进行 softmax 操作。

import torch

tensor = torch.tensor([[[1.0, 2.0, 3.0, 4.0],
                        [2.0, 3.0, 4.0, 1.0],
                        [3.0, 1.0, 2.0, 4.0]],
                       [[4.0, 3.0, 2.0, 1.0],
                        [3.0, 2.0, 1.0, 4.0],
                        [2.0, 4.0, 3.0, 1.0]]])

softmax_tensor = torch.softmax(tensor, dim=2)

print(softmax_tensor)

运行以上代码，得到的输出为：

tensor([[[0.0321, 0.0871, 0.2369, 0.6439],
         [0.1231, 0.3333, 0.0900, 0.4536],
         [0.2689, 0.0180, 0.0498, 0.6633]],

        [[0.6439, 0.2369, 0.0871, 0.0321],
         [0.4536, 0.0900, 0.3333, 0.1231],
         [0.6633, 0.0498, 0.0180, 0.2689]]])