Tensor列表索引本质

wenjurongyu2024-07-22 1:34

张量列表索引本质

单列表索引

将原张量shape对应位置数字更改为索引列表的形状

python 复制代码

data = torch.randn((9, 512, 30))
index_list = [0, 1, 3]
print(data[:, index_list].shape)
'''
[[9, 3, 30]]
'''


# torch.Size([1, 6, 5])
index2_list = torch.IntTensor([[[4, 5, 6, 7, 8],
                               [4, 5, 6, 7, 8],
                               [4, 5, 6, 7, 8],
                               [4, 5, 6, 7, 8],
                               [4, 5, 6, 7, 8],
                               [4, 5, 6, 7, 8]]])
print(data[:, index2_list].shape)
# torch.Size([6, 1, 6, 5, 30])

多列表索引

双列表索引本质就是索引列表之间通过广播机制，达成一致。

python 复制代码

data = torch.randn((6, 512, 30))
# index1: torch.Size([1, 6, 5])
index1 = torch.IntTensor([[[4, 5, 6, 7, 8],
                           [4, 5, 6, 7, 8],
                           [4, 5, 6, 7, 8],
                           [4, 5, 6, 7, 8],
                           [4, 5, 6, 7, 8],
                           [4, 5, 6, 7, 8]]])

# index2: torch.Size([5])
index2 = torch.IntTensor([0, 1, 2, 3, 4])


# 双列表索引
print(data[index1, index2].size())
'''
torch.Size([1, 6, 5, 30])
'''

# 三列表索引
print(data[index1, index2, index1].size())
'''
torch.Size([1, 6, 5])
'''

列表索引本质

多列表是先广播。

如果是单列表：

$1\] 返回：\[data\[1\]$
$1, 2, 3\]，返回：\[data\[1\], data\[2\], data\[3\]$
$1, 2, 3\]\]，返回： \[ \[ data\[1\], data\[2\], data\[3\]$

如果是双列表：

$1\], \[2\] : 返回：\[ data\[1, 2\]$
$1\]\], \[1, 2\] : 先广播：\[\[1, 1\]\], \[\[1, 2\]\] 重叠 \[\[(1,1), (1, 2)\]\] 返回： \[ \[ data\[1, 1\], data\[1,2\]$

多列表以此类推。