pytorch中利用gather函数取出标签的预测概率的理解

背景

在pytorch中神经网络训练输出为one-hot编码，假设y_h为2个样本，在3个类别的情况下的输出：

y_h= torch.tensor([[0.1, 0.2, 0.7], [0.2, 0.2, 0.6]])

放在表格中便于理解：(此处类别从0开始编号)

样本序号	类别0预测概率	类别1预测概率	类别2预测概率
样本1	0.1	0.2	0.7
样本2	0.2	0.2	0.6

y存放样本对应的标签类别：

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

即：第一个样本属于第0类，第二个样本属于第2类
则提取预测概率：

out=y_h.gather(1, y.view(-1, 1))

输出为：

tensor([[0.1000],
		[0.6000]])

这样就把预测输出中标签的预测概率提取出来了。

理解

在pytorch中，函数

torch.gather(input, dim, index, out=None) → Tensor

沿给定轴 dim ,将输入索引张量 index 指定位置的值进行聚合.
对一个 2 维张量,输出可以定义为:

out[i][j] = input[ index[i][j] ][j] # if dim == 0
out[i][j] = input[i][ index[i][j] ]  # if dim == 1

在这个例子中，

index=y.view(-1, 1)=tensor([[0],
						    [2]])

则输出为

out=y_h[i][index]=tensor([[0.1000],
					   	[0.6000]])

进而实现标签的预测概率提取。
需要注意的是index需要是LongTensor类型