PyTorch torch.index_add-Funktion

Pytorch torch 参考手册PyTorch torch Referenzhandbuch

torch.index_addist eine Funktion in PyTorch, mit der die Werte eines Quelltensors an angegebene Indexpositionen addiert werden. Entlang der angegebenen Dimensiondim, beiindexan den angegebenen Indexpositionen hinzufügensourcedie Werte von.

Funktionsdefinition

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

Parameters:

  • input(Tensor): Eingabe-Tensor.
  • dim(int): Die Dimension des Index.
  • index(Tensor): Ein eindimensionaler Integer-Tensor, der die Positionen angibt, an denen hinzugefügt werden soll.
  • source(Tensor): Quell-Tensor, die hinzuzufügenden Werte.
  • alpha(float, optional): Skalierungsfaktor für source, Standardwert ist 1.

Rückgabewert:

  • torch.Tensor: Gibt den geänderten Tensor zurück.

Anwendungsbeispiele

Beispiel

import torch

# Eingabe-Tensor erstellen
input = torch.randn(4, 5)

# Index und source erstellen
index = torch.tensor([0, 2, 3])
source = torch.randn(3, 5)

# Entlang dim=0 addieren
output = torch.index_add(input, dim=0, index=index, source=source)

print("Eingabeform:", input.shape)
print("Index:", index)
print("Quellform:", source.shape)
print("Ergebnisform:", output.shape)
print("nErgebnis:")
print(output)

Ausgabeergebnis:

输入形状: 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]])

Beispiel

import torch

# Alpha-Parameter verwenden, um source zu skalieren
input = torch.zeros(5)
index = torch.tensor([0, 2, 4])
source = torch.tensor([10, 20, 30])

# alpha=2 bedeutet, nach Multiplikation von source mit 2 zu addieren
output = torch.index_add(input, dim=0, index=index, source=source, alpha=2)

print("Eingabe:", input)
print("Quelle:", source)
print("Ergebnis nach alpha=2:", output)

Ausgabeergebnis:

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

Beispiel

import torch

# Entlang einer anderen Dimension addieren
input = torch.zeros(3, 4, 5)
index = torch.tensor([1, 3])
source = torch.randn(2, 4, 5)

# Entlang dim=1 addieren
output = torch.index_add(input, dim=1, index=index, source=source)

print("Eingabeform:", input.shape)
print("Index-Form:", index.shape)
print("Quellform:", source.shape)
print("Ergebnisform:", output.shape)

Ausgabeergebnis:

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

Beispiel

import torch

# Anwendung in neuronalen Netzen: Attention-Mechanismus
# Angenommen, es gibt mehrere Schlüssel-Wert-Paare, die zur Abfrage aggregiert werden müssen

# Simuliere Abfrage und Schlüssel-Wert
num_queries = 2
num_kv = 4
dim = 3

# Abfrageindizes
query_idx = torch.tensor([0, 1])
# Zugehörige Werte
values = torch.randn(num_queries, dim) * 10

# Ausgabe
output = torch.zeros(num_kv, dim)
# Wert an der entsprechenden Position hinzufügen
output = torch.index_add(output, dim=0, index=query_idx, source=values)

print("Abfrageindizes:", query_idx)
print("Werte:", values)
print("Aggregationsergebnis:", output)

Hinweis:torch.index_addändert den ursprünglichen Eingabe-Tensor nicht, sondern gibt einen neuen Tensor zurück. Wennindexgibt es doppelte Indizes in [index], werden die Werte akkumuliert.alphaDer [alpha]-Parameter kann verwendet werden, um die Quellwerte zu skalieren.

Pytorch torch 参考手册PyTorch torch Referenzhandbuch

Andere Erweiterungen