如何用1D张量对PyTorch中的3D张量进行切片？

针对用[Batch]形状的1D张量idx，对[Batch, Sequence, Dim]形状的3D张量t按Sequence维度切片的需求，最高效的实现方式是利用PyTorch的高级索引，完全不需要one-hot编码、乘法求和这类冗余操作。

问题场景回顾 #

给定输入：

import torch
from torch import nn

t = torch.arange(24).view(2,3,4)
# 输出：
# tensor([[[ 0,  1,  2,  3],
#          [ 4,  5,  6,  7],
#          [ 8,  9, 10, 11]],
# 
#         [[12, 13, 14, 15],
#          [16, 17, 18, 19],
#          [20, 21, 22, 23]]])

idx = torch.tensor([0,2])
# 输出：tensor([0, 2])

期望得到每个batch对应idx位置的Sequence切片：tensor([[ 0, 1, 2, 3], [20, 21, 22, 23]])

高效解法代码 #

直接使用高级索引定位：

result = t[torch.arange(t.shape[0]), idx]
print(result)
# 输出：
# tensor([[ 0,  1,  2,  3],
#         [20, 21, 22, 23]])

原理说明 #

• torch.arange(t.shape[0])生成从0到Batch-1的索引张量，对应每个batch的位置；
• 将其与idx组合作为索引，PyTorch会自动为每个batch选取idx指定的Sequence维度切片，直接得到目标结果；
• 这种方式直接通过索引定位目标元素，无额外中间张量生成，计算效率和内存占用远优于原解法。

原解法的问题 #

你提到的one-hot+乘法求和方式：

one_hot_idx = nn.functional.one_hot(idx.long(), num_classes=t.shape[1]).unsqueeze(-1)
(t*one_hot_idx).sum(1)

会先生成和t同形状的one-hot张量，再做逐元素乘法，最后求和，不仅计算量更大，还会占用额外内存存储中间张量，当Batch或Sequence维度较大时，性能差距会非常明显。

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