PyTorch中3D张量被双3D张量索引的机制及维度疑问解析

我有一个3D张量x_k，还有两个作为索引的3D张量h_att和i_att，想搞清楚q = x_k[h_att,i_att]的具体工作机制。相关代码和张量信息如下：

import torch
import numpy as np
import torch.nn as nn

x_k = torch.arange(1,25).reshape(2,4,3)

x_k的输出：

tensor([[[ 1,  2,  3],
         [ 4,  5,  6],
         [ 7,  8,  9],
         [10, 11, 12]],

        [[13, 14, 15],
         [16, 17, 18],
         [19, 20, 21],
         [22, 23, 24]]])

i_att的定义与输出：

i_att = torch.tensor([[[0,1],[1,2],[2,3],[3,0]]])

tensor([[[0, 1],
         [1, 2],
         [2, 3],
         [3, 0]]])

h_att的定义与输出：

h_att = torch.arange(2).reshape(2, 1, 1).long()

tensor([[[0]],

        [[1]]])

执行q = x_k[h_att,i_att]后，q的输出：

tensor([[[[ 1,  2,  3],
          [ 4,  5,  6]],

         [[ 4,  5,  6],
          [ 7,  8,  9]],

         [[ 7,  8,  9],
          [10, 11, 12]],

         [[10, 11, 12],
          [ 1,  2,  3]]],


        [[[13, 14, 15],
          [16, 17, 18]],

         [[16, 17, 18],
          [19, 20, 21]],

         [[19, 20, 21],
          [22, 23, 24]],

         [[22, 23, 24],
          [13, 14, 15]]]])

疑问解答 #

1. 为什么输出张量的维度是(2,4,2,3)？ #

首先明确x_k的维度是(2,4,3)，对应(batch_size, 序列长度, 特征维度)。

当用两个索引张量h_att（维度(2,1,1)）和i_att（维度(1,4,2)）索引x_k时，PyTorch会先对二者执行广播（broadcasting），让它们的维度完全匹配：

• h_att会被广播为(2,4,2)，即每个batch下的4个序列位置、每个位置的2个索引，都对应同一个batch编号（第0个batch全是0，第1个全是1）
• i_att会被广播为(2,4,2)，即两个batch共享同一套序列索引规则

广播完成后，每个索引位置(b, s, idx)对应的是x_k[b, i_att[b,s,idx], :]，这是一个长度为3的特征向量。最终把所有结果按索引维度拼接，就得到了(2,4,2,3)的输出：

• 第0维：batch数2
• 第1维：序列长度4
• 第2维：每个序列位置选取的2个样本
• 第3维：特征维度3

2. 将i_att改为2D张量后结果一致的原因 #

修改后的i_att是(4,2)的2D张量，当和h_att（(2,1,1)）一起索引时，依然触发广播机制：

• h_att会被广播到(2,4,2)
• 2D的i_att会自动添加一个前置维度变为(1,4,2)，再广播到(2,4,2)

广播后的两个索引张量和原场景完全相同，因此取到的元素也完全一致，输出结果自然没有变化。

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