pytorch 计算topk_PyTorch常用函数：torch.topk; torch.index_select; torch.lt

torch.topk

torch.topk(input, k, dim=None, largest=True, sorted=True, out=None) -> (Tensor, LongTensor)

pytorch中文官网文档：http://www.mamicode.com/info-detail-2217311.html

沿给定dim维度返回输入张量input中 k 个最大值。

如果不指定dim，则默认为input的最后一维。

如果为largest为 False ，则返回最小的 k 个值。

返回一个元组 (values,indices)，其中indices是原始输入张量input中测元素下标。

如果设定布尔值sorted 为_True_，将会确保返回的 k 个值被排序。

参数:

input (Tensor) – 输入张量

k (int) – “top-k”中的k

dim (int, optional) – 排序的维

largest (bool, optional) – 布尔值，控制返回最大或最小值

sorted (bool, optional) – 布尔值，控制返回值是否排序

out (tuple, optional) – 可选输出张量 (Tensor, LongTensor) output buffer

------------------------------------实例------------------------------------

假设神经网络的输出如下，为二分类。batch_size=4

import torch

output = torch.tensor([[-5.4783, 0.2298],

[-4.2573, -0.4794],

[-0.1070, -5.1511],

[-0.1785, -4.3339]])

得到其top1值操作如下：

maxk = max((1,)) # 取top1准确率，若取top1和top5准确率改为max((1,5))

_, pred = output.topk(maxk, 1, True, True)

topk参数中，maxk取得是top1准确率，dim=1是按行取值， largest=1是取最大值

结果如下，

_

tensor([[ 0.2298],

[-0.4794],

[-0.1070],

[-0.1785]])

pred

tensor([[1],

[1],

[0],

[0]])

_是top1的值，pred是最大值的索引(size=4*1)，一般会进行转置处理同真实值对比

index_select

anchor_w = self.FloatTensor(self.scaled_anchors).index_select(1, self.LongTensor([0]))

参数说明：index_select(x, 1, indices)

1代表维度1，即列，indices是筛选的索引序号。

例子：

import torch

x = torch.linspace(1, 12, steps=12).view(3,4)

print(x)

indices = torch.LongTensor([0, 2])

y = torch.index_select(x, 0, indices)

print(y)

z = torch.index_select(x, 1, indices)

print(z)

z = torch.index_select(y, 1, indices)

print(z)

结果：

tensor([[  1.,   2.,   3.,   4.],

[  5.,   6.,   7.,   8.],

[  9.,  10.,  11.,  12.]])

tensor([[  1.,   2.,   3.,   4.],

[  9.,  10.,  11.,  12.]])

tensor([[  1.,   3.],

[  5.,   7.],

[  9.,  11.]])

tensor([[  1.,   3.],

[  9.,  11.]])

torch.lt

torch.lt(input, other, out=None) → Tensor

Computes \text{input} < \text{other}input

The second argument can be a number or a tensor whose shape is broadcastable with the first argument.

Parametersinput (Tensor) – the tensor to compare

other (Tensororpython:float) – the tensor or value to compare

out (Tensor,optional) – the output tensor that must be a BoolTensorReturnsA torch.BoolTensor containing a True at each location where comparison is trueReturn type

Example: