本文介绍如何在 pytorch 中对二维张量的每一行，根据动态指定的模式（取最大值或最小非零值），高效、向量化地获取对应元素的列索引，避免循环，兼顾可读性与计算性能。
本文介绍如何在 pytorch 中对二维张量的每一行，根据动态指定的模式（取最大值或最小非零值），高效、向量化地获取对应元素的列索引，避免循环，兼顾可读性与计算性能。
在深度学习实践中，常需对 batch 维度下的每条样本（即张量的每一行）执行差异化 top-k 索引选择：例如，某些样本取最大值位置，另一些则取“最小非零值”的位置。这类需求无法直接用 torch.topk 统一处理，但可通过符号控制 + 值掩码实现完全向量化——无需 Python 循环，全程基于张量操作，保证 GPU 友好性与训练吞吐。
核心思路是：将两类目标统一为“带符号加权后的 top-1 选取”。具体步骤如下：
- 定义选择策略向量 k：k[i] == 1 表示第 i 行取最大值；k[i] == 0 表示取最小非零值；
- 构造符号掩码：将 k 转换为 +1（保持原序）或 -1（反转顺序），即 k_sign = k + (-1) * (k == 0).float()；
- 符号加权：x_signed = x * k_sign.unsqueeze(1)，使“求最小”等价于“对负值求最大”；
- 屏蔽零值干扰：使用 masked_fill(x == 0, -inf) 将所有 0 替换为负无穷，确保其在 topk 中永远被排除；
- 统一 topk 提取：调用 x_filled.topk(1, dim=1) 获取每行 top-1 的索引，并 squeeze() 去除冗余维度。
以下是完整可运行示例代码：
import torch
# 输入张量：shape = (batch_size, N)
x = torch.tensor([[4, 3, 1, 4, 2],
[0, 0, 2, 3, 4],
[4, 4, 3, 0, 3]]).float() # 注意：必须为 float 类型，支持 -inf 填充
# 动态策略向量：1→取最大值，0→取最小非零值
k = torch.tensor([1, 0, 0]).long()
# 步骤1：生成符号向量（1 → +1；0 → -1）
k_sign = k + (-1.0) * (k == 0).float()
# 步骤2：按行应用符号（广播至列维度）
x_signed = x * k_sign.unsqueeze(1)
# 步骤3：将原始张量中为 0 的位置设为 -inf（确保不被选中）
x_filled = x_signed.masked_fill(x == 0, float('-inf'))
# 步骤4：每行取 top-1 索引（返回值、索引；我们只需索引）
_, indices = x_filled.topk(1, dim=1)
output = indices.squeeze(-1) # shape: (batch_size,)
print("输出索引:", output) # tensor([0, 2, 2])
print("对应值:", x.gather(1, output.unsqueeze(1)).squeeze()) # tensor([4., 2., 3.])
✅ 关键注意事项：
- 输入张量 x 必须为浮点类型（如 float32），否则 masked_fill(..., -inf) 会报错；
- k_sign.unsqueeze(1) 是实现行级广播的关键，确保符号正确作用于整行；
- masked_fill 仅修改用于排序的临时张量，原始 x 不受影响；
- 若某行全为零（无非零元素），该行将全部为 -inf，topk 仍返回首个索引（但值为 -inf），建议业务层额外校验；
- 扩展性良好：将 topk(1, ...) 改为 topk(k_val, ...) 即可支持每行不同数量的 top-k 索引（需 k 为标量或适配广播）。
该方案兼具简洁性与工程鲁棒性，适用于动态采样、课程学习（curriculum learning）、稀疏注意力掩码构建等场景，是 PyTorch 高级索引技巧的典型实践。
