PyTorch中如何根据类别张量索引对应向量生成目标张量？

你需要的功能可以通过torch.gather或者PyTorch高级索引实现，torch.scatter是反向的散射操作，确实不匹配你的需求。

实现逻辑 #

核心是按照A张量每个位置的类别索引，从B张量对应batch的两个类别向量中取出对应向量，填充到输出的对应空间位置。

代码示例 #

方法1：使用高级索引（更直观） #

import torch

# 定义张量维度参数
B, H, W, D = 2, 3, 4, 8

# 构造测试输入
A = torch.randint(0, 2, (B, H, W)) # BxHxW的类别索引
B_tensor = torch.randn(B, 2, D) # Bx2xD的类别向量，此处避免和变量名B冲突重命名为B_tensor

# 实现索引取值 + 维度调整
result = B_tensor[torch.arange(B, device=A.device)[:, None, None], A].permute(0, 3, 1, 2)

# 验证输出形状，结果为 torch.Size([2, 8, 3, 4])，符合BxDxHxW要求
print(result.shape)

方法2：使用torch.gather（更适合多维度扩展场景） #

# 调整B的维度顺序为 BxDx2
B_reshaped = B_tensor.permute(0, 2, 1)
# 扩展A的维度为BxDxHxW，作为索引（用expand无需额外复制内存）
A_index = A.unsqueeze(1).expand(-1, D, -1, -1)
# 沿最后一个维度按索引取值
result = B_reshaped.gather(dim=2, index=A_index)

print(result.shape)

两种方法的输出结果完全一致，你可以根据自己的使用场景选择。

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