本文系统对比了在PyTorch中对2D二值张量按行提取1所在列索引的多种方法，涵盖torch.nonzero、nonzero_static、索引广播、vmap等策略，并基于CPU/GPU实测性能给出明确推荐。
本文系统对比了在pytorch中对2d二值张量按行提取`1`所在列索引的多种方法，涵盖`torch.nonzero`、`nonzero_static`、索引广播、`vmap`等策略，并基于cpu/gpu实测性能给出明确推荐。
在深度学习与图神经网络等场景中，常需对形如 [m, n] 的二值张量（仅含0/1）按行提取所有值为 1 的列索引——即为每行生成一个动态长度的索引列表。直观做法是调用 torch.nonzero()，但其逐元素扫描+动态内存分配的机制常被质疑为性能瓶颈；而较新的 torch.nonzero_static(size=...) 虽支持预分配输出空间，却不支持 dim 参数，也无法直接沿指定维度批量执行，且当前尚不支持 CUDA 后端，极大限制了其在GPU训练中的实用性。
针对这一需求，实践中存在以下主流技术路径，我们结合原理与实测数据逐一分析：
✅ 推荐方案：索引广播 + 按行排序（最高效、简洁、GPU友好）
该方法无需调用任何索引查找函数，完全利用向量化操作：
# 假设 data.shape == [m, n]，dtype=torch.bool 或 torch.long index_array = torch.arange(n, device=data.device).unsqueeze(0) # shape: [1, n] masked_indices = data * index_array # 自动广播，0位置保持0，1位置填入列索引 _, sorted_indices = masked_indices.sort(dim=1, descending=True) # 每行内降序排列，有效索引靠前
结果 sorted_indices 中，每行前 k 个元素即为该行所有 1 对应的列索引（k 为该行1的个数），后续填充为 0（因 data*index_array 中0位置乘后仍为0）。若需严格分离有效索引与填充，可配合 torch.cumsum(data, dim=1) 获取每行有效长度，再用 torch.gather 提取——但多数下游任务（如稀疏注意力掩码、邻接索引采样）可直接使用排序后的张量。
✅ 优势：纯张量运算、无Python循环、完美支持CUDA、内存访问连续、实测在CPU/GPU上均显著快于nonzero系列（见下表）。
⚠️ 注意：返回的是固定形状 [m, n] 张量，而非变长列表；若业务强依赖每行独立的1D索引列表（如传入torch.scatter），则需额外后处理。
⚠️ nonzero_static：潜力大但当前受限
尽管 nonzero_static 理论上可通过 view(-1) + 全局索引反解实现按行映射（row = idx // n, col = idx % n），但存在两大硬伤：
- ❌ CUDA 不可用：调用会报错
RuntimeError: nonzero_static is not implemented for CUDA tensors； - ❌
vmap效果差：torch.func.vmap(torch.nonzero_static)触发未实现警告，且实测比朴素nonzero慢2–3倍（见测试结果），因其底层尚未实现批处理规则。
因此，除非你明确处于CPU环境且需严格控制最大输出尺寸，否则不建议优先选用 nonzero_static。
? torch.nonzero()：被低估的“够用”方案
传统观点认为 nonzero() 因动态内存分配而慢，但大规模实测（见文末三组数据）表明：
- 在
[2000, 1000]及更大规模下，nonzero()GPU版本稳定快于所有nonzero_static变体； - 其GPU内核高度优化，实际瓶颈往往不在内存分配，而在后续的
scatter或gather操作； - 若接受“索引对”（
[row_idx, col_idx]）形式输出（即两列张量），可跳过sort和reshape，进一步提速。
? 性能实测关键结论（单位：秒/次，平均100轮）
| 方法 | [200,100] |
[2000,1000] |
[20000,10000] |
|---|---|---|---|
| GPU 索引广播 | 0.00015 | 0.00129 | 0.31346 |
| GPU nonzero | 0.00019 | 0.00198 | 0.30350 |
| CPU 索引广播 | 0.00033 | 0.00466 | 0.55859 |
| CPU nonzero | 0.00051 | 0.00575 | 0.67861 |
| nonzero_static（CPU） | 0.00035 | 0.01028 | 1.10534 |
| vmap + nonzero_static | 0.00191 | 0.03645 | 2.68011 |
? 总结建议：
- 首选 GPU 索引广播法：代码简洁、性能最优、兼容性最好；
- 次选 GPU
nonzero()：若需天然的(row, col)索引对，且不介意后续 reshape；- 避免
nonzero_static+vmap：当前实现低效，且无CUDA支持；- 慎用CPU方案：除非模型本身完全CPU运行，否则GPU加速收益远超算法微优化。
最终，高效张量编程的核心在于用向量化替代索引查找，用广播替代循环，用GPU原生算子替代自定义逻辑——本例正是这一理念的典型实践。
