torch.softmax()和torch.argmax()

torch.softmax()

用于计算给定输入张量上的 softmax 激活函数。softmax 函数通常用于机器学习中的多类分类问题，其目标是预测属于每个类的输入的概率分布

函数

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

• input（张量） - 输入张量。
• dim（int，可选）- 计算 softmax 函数的维度。默认值为 -1，表示最后一个维度。
• dtype（torch.dtype，可选）：输出张量的数据类型

torch.softmax()例子

import torch
 
# create a tensor with shape (3, 4)
x = torch.randn(3, 4)
 
# compute the softmax of the tensor along the last dimension
y = torch.softmax(x, dim=-1)
 
# print the original and softmaxed tensors
print(x)
print(y)

tensor([[-0.6228, -1.5915, -0.1148, -0.3931],
        [-1.4639,  0.7532, -0.4272,  1.3508],
        [ 0.4406, -1.3459, -1.1276, -0.1037]])
tensor([[0.2018, 0.0877, 0.3189, 0.3916],
        [0.0481, 0.1923, 0.0734, 0.6862],
        [0.4705, 0.1051, 0.1253, 0.2991]])

torch.argmax()

用于沿指定维度查找张量中最大值的索引。它以张量的形式返回最大值的索引。

torch.argmax(input, dim=None, keepdim=False)

• input（张量） - 输入张量。
• dim（int，可选）- 用于查找最大值的维度。如果未指定，该函数将返回平展最大值的索引。
• keepdim（布尔值，可选） - 是否保留输入张量的维度。默认值为 False

例子1

import torch
 
# create a tensor with shape (3, 4)
x = torch.tensor([
[1, 2, 3, 4], 
[5, 6, 8, 7], 
[19, 10, 11, 19]
])
 
# find the index of the maximum value along the last dimension
y = torch.argmax(x, dim=-1)
#请注意，如果一行中有多个最大值，则该函数将返回最大值第一次出现的索引
# print the original tensor and the indices of the maximum values
print(x)
print(y)

tensor([[ 1,  2,  3,  4],
        [ 5,  6,  8,  7],
        [ 19, 10, 11, 19]])
tensor([3, 2, 0])

创建一个带有形状的张量，并使用该函数找到沿最后一个维度的最大值的索引。生成的张量具有 的形状，并包含输入张量每一行沿最后一个维度的最大值的索引。

111请注意，如果一行中有多个最大值，则该函数将返回最大值第一次出现的索引。