PyTorch中如何利用二维索引张量为目标张量批量赋值（非循环方式）

我之前处理过类似的批量索引赋值问题，循环实现不仅代码繁琐，在数据量较大时还会拖慢运行速度——用矢量化的高级索引就能完美替代循环，效率直接拉满！下面针对PyTorch和NumPy两种常用场景分别给出实现方案：

PyTorch 实现方案 #

假设你的索引张量inds形状为[B,1,N,2]，目标张量target形状为[B,1,H,W]，核心思路是拆分索引维度，利用广播机制生成批量索引，再通过高级索引直接赋值：

import torch

# 先获取批次大小B
B = inds.size(0)

# 1. 拆分行、列索引：从[B,1,N,2]提取出[B,N]的行/列索引
h_indices = inds[:, 0, :, 0]  # 对应目标张量的H维度索引
w_indices = inds[:, 0, :, 1]  # 对应目标张量的W维度索引

# 2. 生成批量索引：每个批次的N个点都对应自身批次ID，形状广播为[B,N]
batch_indices = torch.arange(B).unsqueeze(1).expand(-1, N)

# 3. 一次性完成所有位置赋值
target[batch_indices, 0, h_indices, w_indices] = 1

更简洁的写法 #

你可以省略显式生成batch_indices的步骤，利用PyTorch的索引广播特性直接实现：

target[torch.arange(B).unsqueeze(1), 0, inds[:,0,:,0], inds[:,0,:,1]] = 1

NumPy 实现方案 #

如果是用NumPy数组，思路完全一致，只是API略有不同：

import numpy as np

B = inds.shape[0]
h_indices = inds[:, 0, :, 0]
w_indices = inds[:, 0, :, 1]

# 生成批量索引，用np.newaxis实现维度扩展
batch_indices = np.arange(B)[:, np.newaxis]

# 赋值操作
target[batch_indices, 0, h_indices, w_indices] = 1

注意事项 #

• 确保你的索引值在合法范围内：h_indices必须在[0, H-1]之间，w_indices必须在[0, W-1]之间，否则会触发索引越界错误。
• 这种矢量化方法的效率远高于循环，尤其是当批次大小B和点数量N较大时，性能提升非常显著。

内容的提问来源于stack exchange，提问作者loran d