pytorch花式索引提取topk的張量

問題設定

或者說，有一個(bs, dim, L)的大張量，索引的index形狀為(bs, X)，想得到一個(bs, dim, X)的reduced向量。我們在進行topk操作（以減少計算量）的時候經常碰到這種情況。
給出如下兩種實現方法，分別使用花式索引（參考informer的代碼）以及pytorch的gather方法

代碼實現

索引方法

參考https://blog.csdn.net/qq_36560894/article/details/122005808

feature = torch.rand(2,16,4*4)
indices = torch.randint(0,16, (2, 3))
indices
indices_expand = indices.unsqueeze(1).expand(-1, dim, -1).to(torch.long) # (bs, dim, H*W)
indices_expand.shape
indices_expand[:,1,:] # 結果和indices一致，說明在第二個channel上，每個樣本的索引是一樣的
bs,dim=feature.shape[:2]
bs,dim 
feature_reduce = feature.view(bs, dim, -1)[torch.arange(bs)[:, None, None], torch.arange(dim)[None,:,None], indices_expand]
feature_reduce.shape

gather方法

reduce_feature = torch.gather(feature, 2, indices_expand)

驗證

兩種方法得到的結果完全相同

補充知識

expand方法

在 PyTorch 中，expand() 方法用于擴展張量的大小。它會在不實際復制數據的情況下，重復張量的元素以填充新的形狀。這個方法可以用于廣播操作，以便在執行一些需要相同形狀的張量之間的數學運算時，使它們具有相同的形狀。

下面是使用 expand() 方法的基本用法：

import torch

# 創建一個原始張量
x = torch.tensor([[1, 2, 3],
                  [4, 5, 6]])

# 使用 expand 擴展張量的大小
expanded_x = x.expand(2, 3, 4)  # 擴展成維度為(2, 3, 4)的張量

print(expanded_x)

在上面的例子中，我們首先創建了一個形狀為 (2, 3) 的原始張量 x。然后，我們使用 expand() 方法將其擴展成一個維度為 (2, 3, 4) 的新張量 expanded_x，該張量的形狀是在原始張量形狀的基礎上每個維度都擴展了一倍。

需要注意的是，expand() 方法只能用于增加張量的大小，不能減小。另外，擴展后的張量與原始張量共享底層數據，因此在原始張量上進行的任何修改都會反映在擴展后的張量上，反之亦然。

gather方法

在 PyTorch 中，gather() 方法用于從輸入張量中按照指定索引提取元素。這個方法通常用于根據索引收集特定的元素，例如根據類別索引從分類得分張量中獲取對應類別的得分。

下面是使用 gather() 方法的基本用法：

import torch

# 創建一個輸入張量
input_tensor = torch.tensor([[1, 2],
                             [3, 4],
                             [5, 6]])

# 創建一個索引張量
indices = torch.tensor([[0, 0],
                        [1, 0]])

# 使用 gather 方法根據索引收集元素
output_tensor = torch.gather(input_tensor, dim=1, index=indices)

print(output_tensor)

在上面的例子中，我們首先創建了一個形狀為 (3, 2) 的輸入張量 input_tensor，以及一個形狀為 (2, 2) 的索引張量 indices。然后，我們使用 gather() 方法從輸入張量 input_tensor 中按照索引張量 indices 收集元素。

在 gather() 方法中，參數 dim 指定了在哪個維度上進行收集操作，而 index 參數指定了收集元素所使用的索引張量。

需要注意的是，索引張量 indices 的形狀必須與輸出張量的形狀一致，或者是可以廣播成與輸出張量形狀一致的形狀。

randint

torch.randint() 是 PyTorch 中用于生成隨機整數張量的函數。它可以生成一個張量，其中的元素是在指定范圍內隨機抽樣的整數。

下面是 torch.randint() 的基本用法示例：

import torch

# 生成一個形狀為 (3, 3) 的隨機整數張量，范圍是 [0, 10)
random_integers = torch.randint(low=0, high=10, size=(3, 3))

print(random_integers)

在上面的示例中，我們使用了 torch.randint() 函數來生成一個形狀為 (3, 3) 的隨機整數張量，其中的元素取值范圍在閉區間 [low, high) 內，即從 0 到 9。

torch.randint() 函數的主要參數包括：

• low：生成的隨機整數的最小值（包含）。
• high：生成的隨機整數的最大值（不包含）。
• size：生成的張量的形狀。

你也可以不指定 low 參數，默認情況下它為 0。此外，還可以使用其他參數來控制生成的隨機整數張量的設備類型、數據類型等。