PyTorch中如何根据二维张量索引将三维张量对应位置置0

你可以直接用PyTorch原生的多维高级索引完成需求，不需要写循环，执行效率很高。

实现思路 #

对于形状为(B, N, d)的张量P，要定位到需要置零的位置，需要同时给出三个维度的坐标：

• 第一维（批次维）坐标：和Q形状一致，第i行的所有元素值为i，对应该行属于第i个批次
• 第二维（序列维）坐标：和Q形状一致，第j列的所有元素值为j，对应该列属于序列的第j个位置
• 第三维（特征维）坐标：直接使用Q本身的值即可
  三个维度的坐标形状完全匹配，可以直接定位到所有需要置零的元素，一次性完成赋值。

可运行代码示例 #

import torch

# 题目给出的示例输入
P = torch.randn(2, 3, 6)
Q = torch.tensor([[0, 3, 4], [2, 1, 3]])

B, N = Q.shape
# 构造前两个维度的索引张量
b_idx = torch.arange(B).unsqueeze(1).expand(-1, N)
n_idx = torch.arange(N).unsqueeze(0).expand(B, -1)

# 定位到目标位置直接置零
P[b_idx, n_idx, Q] = 0

效果验证 #

执行上述代码后，以下位置会被准确置为0，其余元素保持原值不变：

• P[0, 0, 0] = 0
• P[0, 1, 3] = 0
• P[0, 2, 4] = 0
• P[1, 0, 2] = 0
• P[1, 1, 1] = 0
• P[1, 2, 3] = 0

提示：这种向量化实现没有Python层面的循环开销，在张量维度较大时运行速度远高于逐元素遍历的写法，且没有额外的API依赖，所有稳定版PyTorch都可以直接运行。

内容的提问来源于stack exchange，提问作者Young By