Función PyTorch torch.index_add

Pytorch torch 参考手册Manual de referencia de Pytorch torch

torch.index_addEs una función de PyTorch que se utiliza para sumar los valores del tensor fuente a las posiciones de índice especificadas. A lo largo de la dimensión especificadadim, enindexlas posiciones de índice especificadas se añadensourcelos valores de.

Definición de la función

torch.index_add(input, dim, index, source, *, alpha=1)

Parámetros:

  • input(Tensor): tensor de entrada.
  • dim(int): dimensión del índice.
  • index(Tensor): tensor unidimensional de enteros que especifica las posiciones a las que se añadirá.
  • source(Tensor): tensor fuente, los valores que se añadirán.
  • alpha(float, opcional): factor de escala de source, por defecto 1.

Valor de retorno:

  • torch.Tensor: Devuelve el tensor modificado.

Ejemplo de uso

Ejemplo

import torch

# Crear el tensor de entrada
input = torch.randn(4, 5)

# Crear el índice y la fuente
index = torch.tensor([0, 2, 3])
source = torch.randn(3, 5)

# Añadir a lo largo de dim=0
output = torch.index_add(input, dim=0, index=index, source=source)

print(Forma de entrada:, input.shape)
print(Índice:, index)
print(Forma de la fuente:, source.shape)
print(Forma del resultado:, output.shape)
print(nResultado:)
print(output)

El resultado de salida es:

输入形状: torch.Size([4, 5])
索引: tensor([0, 2, 3])
源形状: torch.Size([3, 5])
结果形状: torch.Size([4, 5])

结果:
tensor([[ 1.8435,  0.3463, -0.1024,  0.5678,  0.1234],
        [-0.5678,  0.8901, -0.2345,  0.6789, -0.1234],
        [ 2.3456,  0.4567,  0.7890, -0.3456,  0.5678],
        [-0.7890,  1.2345,  0.3456, -0.8901,  0.2345]])

Ejemplo

import torch

# Usar el parámetro alpha para escalar la fuente
input = torch.zeros(5)
index = torch.tensor([0, 2, 4])
source = torch.tensor([10, 20, 30])

# alpha=2 significa multiplicar source por 2 y luego añadir
output = torch.index_add(input, dim=0, index=index, source=source, alpha=2)

print(Entrada:, input)
print(Fuente:, source)
print(Resultado con alpha=2:, output)

El resultado de salida es:

输入: tensor([0., 0., 0., 0., 0.])
源: tensor([10., 20., 30.])
alpha=2 后的结果: tensor([20.,  0., 40.,  0., 60.])

Ejemplo

import torch

# Añadir a lo largo de otra dimensión
input = torch.zeros(3, 4, 5)
index = torch.tensor([1, 3])
source = torch.randn(2, 4, 5)

# Añadir a lo largo de dim=1
output = torch.index_add(input, dim=1, index=index, source=source)

print(Forma de entrada:, input.shape)
print(Forma del índice:, index.shape)
print(Forma de la fuente:, source.shape)
print(Forma del resultado:, output.shape)

El resultado de salida es:

输入形状: torch.Size([3, 4, 5])
索引形状: torch.Size([2])
源形状: torch.Size([2, 4, 5])
结果形状: torch.Size([3, 4, 5])

Ejemplo

import torch

# Aplicación en redes neuronales: mecanismo de atención
# Supongamos que hay varios pares key-value que deben agregarse a query

# Simular query y key-value
num_queries = 2
num_kv = 4
dim = 3

# Índice de query
query_idx = torch.tensor([0, 1])
# El value correspondiente
values = torch.randn(num_queries, dim) * 10

# Salida
output = torch.zeros(num_kv, dim)
# Añadir el value a la posición correspondiente
output = torch.index_add(output, dim=0, index=query_idx, source=values)

print(Índice de query:, query_idx)
print("Values:", values)
print(Resultado de la agregación:, output)

Nota:torch.index_addNo modifica el tensor de entrada original, sino que devuelve un nuevo tensor. Siindexhay índices duplicados, los valores se acumularán.alphaEl parámetro se puede utilizar para escalar los valores de la fuente.

Pytorch torch 参考手册Manual de referencia de Pytorch torch

Otras extensiones