Python torch.topk() 函数用法

最新推荐文章于 2024-05-31 09:39:14 发布

sweettea~

最新推荐文章于 2024-05-31 09:39:14 发布

阅读量721

点赞数 2

分类专栏： Python 文章标签： python 深度学习人工智能

本文链接：https://blog.csdn.net/sweet_tea_/article/details/128554578

版权

Python 专栏收录该内容

54 篇文章 5 订阅

订阅专栏

torch.topk(input, k, dim=None, largest=True, sorted=True, out=None)

参数：
input -> 输入tensor
k -> 前k个
dim -> 默认为输入tensor的最后一个维度
sorted -> 是否排序
largest -> False表示返回第k个最小值
示例1：

m=torch.arange(0,10)
print(m.topk(3))

torch.return_types.topk(
values=tensor([9, 8, 7]),
indices=tensor([9, 8, 7]))

示例2：

pred = torch.tensor([[-0.5816, -0.3873, -1.0215, -1.0145,  0.4053],
        [ 0.7265,  1.4164,  1.3443,  1.2035,  1.8823],
        [-0.4451,  0.1673,  1.2590, -2.0757,  1.7255],
        [ 0.2021,  0.3041,  0.1383,  0.3849, -1.6311]])
values, indices = pred.topk(4, dim=0, largest=True, sorted=True)
print(values)
print(indices)

tensor([[ 0.7265,  1.4164,  1.3443,  1.2035,  1.8823],
        [ 0.2021,  0.3041,  1.2590,  0.3849,  1.7255],
        [-0.4451,  0.1673,  0.1383, -1.0145,  0.4053],
        [-0.5816, -0.3873, -1.0215, -2.0757, -1.6311]])
tensor([[1, 1, 1, 1, 1],
        [3, 3, 2, 3, 2],
        [2, 2, 3, 0, 0],
        [0, 0, 0, 2, 3]])