Fonction torch.take_along_dim de PyTorch
Manuel de référence Pytorch torch
torch.take_along_dimIl s'agit d'une fonction dans PyTorch utilisée pour obtenir des éléments à des positions d'index le long d'une dimension spécifiée. Elle prend les valeurs deindicesles indices spécifiés dansdimle long de la dimension à partir deinput.
Définition de fonction
torch.take_along_dim(input, indices, dim)
Parameters:
input(Tensor) : tenseur d'entrée.indices(Tensor) : Tenseur d'index, spécifiant les positions des éléments à extraire. La forme doit être compatible avec input le long de la dimension dim.dim(int) : La dimension sur laquelle opérer.
Valeur de retour:
torch.Tensor: Renvoie un nouveau tenseur composé des éléments extraits par les indices.
Exemple d'utilisation
Exemple
import torch
# Créer un tenseur 2D
x = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
print("Tenseur d'origine :")
print(x)
# Extraire les éléments le long de dim=1
indices = torch.tensor([[0, 1, 2],
[2, 1, 0],
[0, 0, 0]])
y = torch.take_along_dim(x, indices, dim=1)
print("nIndices:")
print(indices)
print("nÉléments extraits le long de dim=1:")
print(y)
# Créer un tenseur 2D
x = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
print("Tenseur d'origine :")
print(x)
# Extraire les éléments le long de dim=1
indices = torch.tensor([[0, 1, 2],
[2, 1, 0],
[0, 0, 0]])
y = torch.take_along_dim(x, indices, dim=1)
print("nIndices:")
print(indices)
print("nÉléments extraits le long de dim=1:")
print(y)
Le résultat de sortie est:
原始张量:
tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
索引:
tensor([[0, 1, 2],
[2, 1, 0],
[0, 0, 0]])
沿 dim=1 取出的元素:
tensor([[1, 2, 3],
[6, 5, 4],
[7, 7, 7]])
Exemple
import torch
# Extraire les éléments le long de dim=0
x = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
# Prendre différentes lignes pour chaque colonne
indices = torch.tensor([[0, 1, 2],
[2, 0, 1],
[1, 2, 0]])
y = torch.take_along_dim(x, indices, dim=0)
print("Tenseur d'origine :")
print(x)
print("nIndices:")
print(indices)
print("nElements extraits le long de dim=0:")
print(y)
# Extraire les éléments le long de dim=0
x = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
# Prendre différentes lignes pour chaque colonne
indices = torch.tensor([[0, 1, 2],
[2, 0, 1],
[1, 2, 0]])
y = torch.take_along_dim(x, indices, dim=0)
print("Tenseur d'origine :")
print(x)
print("nIndices:")
print(indices)
print("nElements extraits le long de dim=0:")
print(y)
Le résultat de sortie est:
原始张量:
tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
索引:
tensor([[0, 1, 2],
[2, 0, 1],
[1, 2, 0]])
沿 dim=0 取出的元素:
tensor([[1, 5, 9],
[7, 2, 6],
[4, 8, 3]])
</p>
<div class="example">
<h2 class="example">实例</h2>
<div class="example_code">
<span style="color: Green;font-weight:bold;">import</span> torch<br />
<br />
<span style="color: #a50"># 在3D张量上使用</span><br />
x <span style="color: Gray;">=</span> torch.<span style="color: #05a;">arange</span><span style="color: Olive;">(</span><span style="color: Maroon;">24</span><span style="color: Olive;">)</span>.<span style="color: #05a;">reshape</span><span style="color: Olive;">(</span><span style="color: Maroon;">2</span><span style="color: Gray;">,</span> <span style="color: Maroon;">3</span><span style="color: Gray;">,</span> <span style="color: Maroon;">4</span><span style="color: Olive;">)</span><br />
<span style="color: Green;font-weight:bold;">print</span><span style="color: Olive;">(</span><span style="color: #a11;">"原始形状:"</span><span style="color: Gray;">,</span> x.<span style="color: #05a;">shape</span><span style="color: Olive;">)</span><br />
<br />
<span style="color: #a50"># 沿 dim=1 取元素</span><br />
indices <span style="color: Gray;">=</span> torch.<span style="color: #05a;">tensor</span><span style="color: Olive;">(</span><span style="color: Olive;">[</span><span style="color: Olive;">[</span><span style="color: Maroon;">0</span><span style="color: Gray;">,</span> <span style="color: Maroon;">1</span><span style="color: Gray;">,</span> <span style="color: Maroon;">2</span><span style="color: Olive;">]</span><span style="color: Gray;">,</span><br />
<span style="color: Olive;">[</span><span style="color: Maroon;">2</span><span style="color: Gray;">,</span> <span style="color: Maroon;">0</span><span style="color: Gray;">,</span> <span style="color: Maroon;">1</span><span style="color: Olive;">]</span><span style="color: Olive;">]</span><span style="color: Olive;">)</span><br />
y <span style="color: Gray;">=</span> torch.<span style="color: #05a;">take_along_dim</span><span style="color: Olive;">(</span>x<span style="color: Gray;">,</span> indices<span style="color: Gray;">,</span> dim<span style="color: Gray;">=</span><span style="color: Maroon;">1</span><span style="color: Olive;">)</span><br />
<br />
<span style="color: Green;font-weight:bold;">print</span><span style="color: Olive;">(</span><span style="color: #a11;">"索引形状:"</span><span style="color: Gray;">,</span> indices.<span style="color: #05a;">shape</span><span style="color: Olive;">)</span><br />
<span style="color: Green;font-weight:bold;">print</span><span style="color: Olive;">(</span><span style="color: #a11;">"结果形状:"</span><span style="color: Gray;">,</span> y.<span style="color: #05a;">shape</span><span style="color: Olive;">)</span><br />
<span style="color: Green;font-weight:bold;">print</span><span style="color: Olive;">(</span><span style="color: #a11;">"n结果:"</span><span style="color: Olive;">)</span><br />
<span style="color: Green;font-weight:bold;">print</span><span style="color: Olive;">(</span>y<span style="color: Olive;">)</span><br />
</div>
</div>
<p>输出结果为：</p>
<pre>
原始形状: torch.Size([2, 3, 4])
索引形状: torch.Size([2, 3])
结果形状: torch.Size([2, 3, 4])
结果:
tensor([[[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]],
[[16, 17, 18, 19],
[12, 13, 14, 15],
[20, 21, 22, 23]]])
Exemple
import torch
# Application : sélection d'éléments spécifiques le long de la dimension batch
# Par exemple, sélectionner des paires clé-valeur spécifiques dans un mécanisme d'attention
batch_size = 2
num_heads = 3
seq_len = 4
head_dim = 5
# Simuler les poids d'attention de la requête
attn_weights = torch.randn(batch_size, num_heads, seq_len)
# Prendre les indices top-k pour chaque tête
k = 2
indices = torch.argsort(attn_weights, dim=-1, descending=True)[..., :k]
print("Forme des indices :", indices.shape)
# Simuler le tenseur de valeurs
value = torch.randn(batch_size, num_heads, seq_len, head_dim)
# Extraire les valeurs correspondantes le long de la dimension seq_len
selected_value = torch.take_along_dim(value, indices.unsqueeze(-1).expand(-1, -1, -1, head_dim), dim=2)
print("Forme de la valeur :", value.shape)
print("Forme de la valeur sélectionnée :", selected_value.shape)
# Application : sélection d'éléments spécifiques le long de la dimension batch
# Par exemple, sélectionner des paires clé-valeur spécifiques dans un mécanisme d'attention
batch_size = 2
num_heads = 3
seq_len = 4
head_dim = 5
# Simuler les poids d'attention de la requête
attn_weights = torch.randn(batch_size, num_heads, seq_len)
# Prendre les indices top-k pour chaque tête
k = 2
indices = torch.argsort(attn_weights, dim=-1, descending=True)[..., :k]
print("Forme des indices :", indices.shape)
# Simuler le tenseur de valeurs
value = torch.randn(batch_size, num_heads, seq_len, head_dim)
# Extraire les valeurs correspondantes le long de la dimension seq_len
selected_value = torch.take_along_dim(value, indices.unsqueeze(-1).expand(-1, -1, -1, head_dim), dim=2)
print("Forme de la valeur :", value.shape)
print("Forme de la valeur sélectionnée :", selected_value.shape)
Le résultat de sortie est:
索引形状: torch.Size([2, 3, 4]) Value形状: torch.Size([2, 3, 4, 5]) 选择的Value形状: torch.Size([2, 3, 2, 5])
Note:
torch.take_along_dimIt allows indexing by dimension, which is more flexible thantorch.takebecause the latter always treats the tensor as one-dimensional.