【pyTorch】关于PyTorch的高级索引机制理解 - 指南

y_hat[[0, 1], y]

看起来简短，其实包含了 PyTorch 的高级索引（advanced indexing） 机制，
常用于分类任务中，从模型输出中取出正确类别对应的预测概率（或分数）。

我们来一步步拆解

---

一、先看变量定义

y = torch.tensor([0, 2])
y_hat = torch.tensor([
[0.1, 0.3, 0.6],
[0.3, 0.2, 0.5]
])

可以理解为：

样本编号	模型预测概率（3类）	正确类别
第 0 个样本	[0.1, 0.3, 0.6]	0
第 1 个样本	[0.3, 0.2, 0.5]	2

---

二、代码含义

y_hat[[0, 1], y]

这是一个行列双索引操作，等价于：

torch.tensor([y_hat[0, y[0]], y_hat[1, y[1]]])

也就是：

= [y_hat[0, 0], y_hat[1, 2]]

---

三、一步步替换求值

从矩阵中取出对应元素：

索引	取出的元素	值
y_hat[0, 0]	第 0 行第 0 列	0.1
y_hat[1, 2]	第 1 行第 2 列	0.5

所以：

y_hat[[0, 1], y] = tensor([0.1, 0.5])

---

四、这个操作在深度学习中的用途

这个技巧非常常见！尤其是在 交叉熵损失 (cross-entropy) 计算中。

假设 y_hat 是每个样本对所有类别的预测概率分布，y 是正确标签，
那这行代码可以提取出：

每个样本 “对应正确类别” 的预测概率。

例如：

y_hat[[0, 1], y]
# → tensor([预测第0个样本属于真实类别的概率, 
#           预测第1个样本属于真实类别的概率])

然后就可以算交叉熵：

loss = -torch.log(y_hat[[0, 1], y])

---

✅ 五、总结一句话

代码	含义
y_hat[[0, 1], y]	按行索引 [0, 1] 和对应列索引 y 同时取元素
= [y_hat[0, y[0]], y_hat[1, y[1]]]	等价的展开写法
结果	取出每个样本的正确类别预测值（常用于分类任务）

---

口诀记忆：

y_hat[行索引, 列索引]
会“对齐”地从每一行中取出对应列的元素。
常用来拿到“每个样本在真实类别上的预测概率”。

---

具体理解

数据回顾

y = torch.tensor([0, 2])
y_hat = torch.tensor([
[0.1, 0.3, 0.6],
[0.3, 0.2, 0.5]
])

---

图示说明

          类别0    类别1    类别2
样本0 →   0.1      0.3      0.6
样本1 →   0.3      0.2      0.5

以及标签：

y = [0, 2]

表示：

• 样本 0 的真实类别是 0
• 样本 1 的真实类别是 2

---

执行这句代码：

y_hat[[0, 1], y]

等价于：

取第 0 行的第 y[0]=0 列 → 0.1
取第 1 行的第 y[1]=2 列 → 0.5

---

✅ 可视化标注结果

          类别0    类别1    类别2
样本0 →  [0.1]*    0.3      0.6
样本1 →   0.3      0.2    [0.5]*
星号 * 表示被选中的元素

最终输出：

tensor([0.1, 0.5])

---

用途回顾

在分类任务里：

loss = -torch.log(y_hat[[0, 1], y])

就是取出模型对真实标签类别的预测概率，再取负对数计算交叉熵损失。