如何对形状为(NB,N,2,2)的tensor按batch级索引实现向量化取值？

解决方法 #

你要实现的是按batch维度配对索引第二维的元素，可以用以下两种完全向量化的实现方案，均无需循环：

方案1：使用torch.gather

该方法直接通过维度匹配完成索引，代码如下：

import torch

# 示例张量构造（你可以替换成自己的A、B）
NB = 3
N = 4
A = torch.randn(NB, N, 2, 2)
B = torch.tensor([1, 3, 0])

# 索引逻辑
# 将B调整为和A维度匹配的形状 (NB, 1, 1, 1)
B_reshaped = B.view(-1, 1, 1, 1)
# 在第1维（也就是N所在的维度）做gather，扩展索引匹配后两维的2x2形状
result = A.gather(dim=1, index=B_reshaped.expand(-1, 1, 2, 2))

# 此时result形状为 (NB, 1, 2, 2)，和你预期的输出格式完全一致

方案2：配对batch索引直接取值

逻辑更直观，直接给每个batch的索引和B中的对应索引配对：

# 生成0到NB-1的batch维度索引
batch_indices = torch.arange(NB)
# 配对取值后补一个维度匹配你要的输出格式
result = A[batch_indices, B].unsqueeze(1)

之前写法错误原因 #

• A[B]是索引A的第0维（batch维度），取第B[i]个batch的所有元素，不符合需求
• A[:, B]是对每个batch都取B中所有索引对应的元素，输出形状为(NB, NB, 2, 2)，不是每个batch只取对应一个元素
• B.unsqueeze(1)后的形状是(NB,1)，维度数量和A不匹配，因此索引报错

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