如何在PyTorch中找到张量的第K个和前K个元素？

PyTorch提供了一个方法 torch.kthvalue() 来找到张量的第K个元素。它返回按升序排序的张量中第K个元素的值以及原始张量中该元素的索引。

torch.topk() 方法用于找到前K个或最大的K个元素。它返回张量中前K个或最大的K个元素。

步骤

• 导入所需的库。在下面的Python示例中，所需的Python库为 torch 。确保您已经安装了它。

• 创建一个PyTorch张量并将其打印出来。

• 计算 torch.kthvalue（input，k） 。它返回两个张量。将这两个张量分配给两个新变量 "value" 和 "index" 。这里，input是一个张量，k是一个整数。

• 计算 torch.topk（input，k） 。它返回两个张量。第一个张量具有前K个最大元素的值，第二个张量具有这些元素在原始张量中的索引。将这两个张量分配给新变量 "values" 和 "indices" 。

• 打印张量的第K个元素的值和索引以及张量的前K个元素的值和索引。

示例1

此Python程序显示如何找到张量的第K个元素。

# Python程序查找张量的第K个元素
#导入所需的库
import torch

#创建一个1D张量
T = torch.Tensor([2.334,4.433,-4.33,-0.433,5, 4.443])
print("原始张量：\n", T)

#查找已排序张量中的第3个元素。首先按升序对张量进行排序，然后返回排序后张量中第K个元素的值和元素在原始张量中的索引
value，index = torch.kthvalue(T, 3)

#输出第3个元素的值和索引
print("第3个元素的值：", value)
print("第3个元素的索引：", index)

输出

原始张量：
   tensor([ 2.3340, 4.4330, -4.3300, -0.4330, 5.0000, 4.4430])
第3个元素的值：tensor(2.3340)
第3个元素的索引：tensor(0)

示例2

下面的Python程序显示如何找到张量的前K个或最大的K个元素。

# Python程序查找张量的前K个元素
#导入所需的库
import torch

#创建一个1D张量
T = torch.Tensor([2.334,4.433,-4.33,-0.433,5, 4.443])
print("原始张量：\n", T)

#查找张量中前k=2个或最大的2个元素，返回在原始张量中前2个最大元素的值及它们的索引
values，indices = torch.topk(T, 2)

#输出前2个元素的值和索引
print("前2个元素的值：", values)
print("前2个元素的索引：", indices)

输出

原始张量：
   tensor([ 2.3340, 4.4330, -4.3300, -0.4330, 5.0000, 4.4430])
前2个元素的值：tensor([5.0000, 4.4430])
前2个元素的索引：tensor([4, 5])