Función torch.take_along_dim de PyTorch
Manual de referencia de Pytorch torch
torch.take_along_dimEs una función de PyTorch que se utiliza para obtener elementos en posiciones de índice a lo largo de una dimensión especificada. Segúnindiceslos índices especificados a lo largo dedimla dimensión, de la queinputse toman los valores.
Definición de la función
torch.take_along_dim(input, indices, dim)
Parámetros:
input(Tensor): tensor de entrada.indices(Tensor): tensor de índices, especifica las posiciones de los elementos a obtener. La forma debe ser compatible con input en la dimensión dim.dim(int): la dimensión a lo largo de la cual se opera.
Valor de retorno:
torch.Tensor: devuelve un nuevo tensor formado por los elementos extraídos según los índices.
Ejemplo de uso
Ejemplo
import torch
# Crea un tensor 2D
x = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
print("Tensor original:")
print(x)
# Toma elementos a lo largo 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("nÍndices:")
print(indices)
print("nElementos extraídos a lo largo de dim=1:")
print(y)
# Crea un tensor 2D
x = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
print("Tensor original:")
print(x)
# Toma elementos a lo largo 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("nÍndices:")
print(indices)
print("nElementos extraídos a lo largo de dim=1:")
print(y)
El resultado de salida es:
原始张量:
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]])
Ejemplo
import torch
# Toma elementos a lo largo de dim=0
x = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
# Toma diferentes filas para cada columna
indices = torch.tensor([[0, 1, 2],
[2, 0, 1],
[1, 2, 0]])
y = torch.take_along_dim(x, indices, dim=0)
print("Tensor original:")
print(x)
print("nÍndices:")
print(indices)
print("nElementos extraídos a lo largo de dim=0:")
print(y)
# Toma elementos a lo largo de dim=0
x = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
# Toma diferentes filas para cada columna
indices = torch.tensor([[0, 1, 2],
[2, 0, 1],
[1, 2, 0]])
y = torch.take_along_dim(x, indices, dim=0)
print("Tensor original:")
print(x)
print("nÍndices:")
print(indices)
print("nElementos extraídos a lo largo de dim=0:")
print(y)
El resultado de salida es:
原始张量:
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]]])
Ejemplo
import torch
# Aplicación: selecciona elementos específicos a lo largo de la dimensión del lote
# Por ejemplo, seleccionar key-value específicos en el mecanismo de atención
batch_size = 2
num_heads = 3
seq_len = 4
head_dim = 5
# Simula los pesos de atención del query
attn_weights = torch.randn(batch_size, num_heads, seq_len)
# Toma los índices top-k para cada head
k = 2
indices = torch.argsort(attn_weights, dim=-1, descending=True)[..., :k]
print("Forma de los índices:", indices.shape)
# Simula el tensor de values
value = torch.randn(batch_size, num_heads, seq_len, head_dim)
# Toma el value correspondiente a lo largo de la dimensión seq_len
selected_value = torch.take_along_dim(value, indices.unsqueeze(-1).expand(-1, -1, -1, head_dim), dim=2)
print("Forma de Value:", value.shape)
print("Forma del Value seleccionado:", selected_value.shape)
# Aplicación: selecciona elementos específicos a lo largo de la dimensión del lote
# Por ejemplo, seleccionar key-value específicos en el mecanismo de atención
batch_size = 2
num_heads = 3
seq_len = 4
head_dim = 5
# Simula los pesos de atención del query
attn_weights = torch.randn(batch_size, num_heads, seq_len)
# Toma los índices top-k para cada head
k = 2
indices = torch.argsort(attn_weights, dim=-1, descending=True)[..., :k]
print("Forma de los índices:", indices.shape)
# Simula el tensor de values
value = torch.randn(batch_size, num_heads, seq_len, head_dim)
# Toma el value correspondiente a lo largo de la dimensión seq_len
selected_value = torch.take_along_dim(value, indices.unsqueeze(-1).expand(-1, -1, -1, head_dim), dim=2)
print("Forma de Value:", value.shape)
print("Forma del Value seleccionado:", selected_value.shape)
El resultado de salida es:
索引形状: torch.Size([2, 3, 4]) Value形状: torch.Size([2, 3, 4, 5]) 选择的Value形状: torch.Size([2, 3, 2, 5])
Nota:
torch.take_along_dimPermite indexar por dimensión, lo cual, en comparación contorch.takees más flexible, porque este último siempre trata el tensor como unidimensional.