Select 基於索引的維度上的張量切片

Question

我有一個形狀如下的 PyTorch 張量：(100, 5, 100)。 我需要通過從第二維的每一行中僅選擇一個項目，將其轉換為形狀為 (100, 100) 的張量，這意味着在這 5 個元素中，我只需要一個及其對應的 100 個元素。

要執行此操作，我有第二個形狀為 (100,) 的張量，其索引指定應在每一行中選擇這 5 個項目中的哪一個。

有沒有一種簡單的方法可以執行此選擇而不必過多地弄亂尺寸？

Answer 1

假設張量的索引名為idx並且形狀為(100,) 。 具有稱為source的值的張量。 然后到 select：

result = source[torch.arange(100), idx]