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类型