Fonction torch.scatter_add de PyTorch


Pytorch torch 参考手册Manuel de référence PyTorch torch

torch.scatter_addest une fonction dans PyTorch utilisée pour ajouter les valeurs du tenseur source à des positions spécifiées. Elle vasrcles valeurs selonindexles positions spécifiées auxquelles ajouterinputdans.

Définition de fonction

torch.scatter_add(input, dim, index, src)

Parameters:

  • input(Tensor) : tenseur d'entrée.
  • dim(int) : La dimension à répartir.
  • index(Tensor) : Tenseur d'indices, spécifiant où ajouter les valeurs de src à input.
  • src(Tensor) : Tenseur source, les valeurs à ajouter.

Valeur de retour:

  • torch.Tensor: Renvoie le tenseur modifié.

Exemple d'utilisation

Exemple

import torch

# Créer un tenseur d'entrée
input = torch.zeros(3, 5)

# Créer l'index et la source
index = torch.tensor([[0, 1, 2, 0, 0],
                      [1, 2, 0, 1, 2],
                      [2, 0, 1, 2, 0]])
src = torch.tensor([[1, 1, 1, 1, 1],
                    [2, 2, 2, 2, 2],
                    [3, 3, 3, 3, 3]])

# Répartir et accumuler le long de dim=0
output = torch.scatter_add(input, dim=0, index=index, src=src)

print("Entrée :")
print(input)
print("nIndex:")
print(index)
print("nSource:")
print(src)
print("nRésultat :")
print(output)

Le résultat de sortie est:

输入:
tensor([[0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]])

索引:
tensor([[0, 1, 2, 0, 0],
        [1, 2, 0, 1, 2],
        [2, 0, 1, 2, 0]])

源:
tensor([[1., 1., 1., 1., 1.],
        [2., 2., 2., 2., 2.],
        [3., 3., 3., 3., 3.]])

结果:
tensor([[4., 1., 2., 4., 4.],
        [2., 1., 2., 2., 2.],
        [3., 3., 1., 3., 3.]])

Exemple

import torch

# Utiliser dim=1
input = torch.zeros(3, 5)
index = torch.tensor([[0, 1, 2, 1, 0],
                      [1, 2, 0, 2, 1],
                      [0, 1, 1, 0, 2]])
src = torch.arange(1, 6).float()

output = torch.scatter_add(input, dim=1, index=index, src=src)

print("Scatter le long de dim=1:")
print(output)

Le résultat de sortie est:

沿 dim=1 散布:
tensor([[ 6.,  3.,  3.,  0.,  0.],
        [ 3.,  6.,  3.,  0.,  0.],
        [ 2.,  4.,  5.,  0.,  0.]])

Exemple

import torch

# Scénario d'application pour agréger des valeurs à partir de plusieurs positions
# Par exemple, accumuler les caractéristiques des nœuds voisins dans les réseaux neuronaux de graphes

# Simuler les caractéristiques initiales de 4 nœuds
node_features = torch.zeros(4, 3)

# Simuler les connexions d'arêtes (les nœuds source pointent vers les nœuds cibles)
edge_index = torch.tensor([0, 1, 2, 3, 0, 1])  # Nœuds source des arêtes
edge_weights = torch.tensor([1.0, 2.0, 3.0, 1.5, 2.5, 0.5])

# Créer des valeurs pondérées des caractéristiques des nœuds sources pour chaque arête
src_features = torch.randn(6, 3) * edge_weights.unsqueeze(1)

# Accumuler les caractéristiques vers les nœuds cibles (simplifié ici ; en pratique, cela devrait être basé sur les nœuds cibles des arêtes)
target_nodes = torch.tensor([0, 0, 1, 1, 2, 3])
index = target_nodes

output = torch.scatter_add(node_features, 0, index.unsqueeze(1).expand_as(src_features), src_features)

print("Forme des caractéristiques des nœuds :", node_features.shape)
print("Caractéristiques accumulées :", output)

Remarque :torch.scatter_addne modifie pas le tenseur d'entrée d'origine, mais renvoie un nouveau tenseur. Plusieurs indices peuvent pointer vers la même position, et les valeurs seront accumulées. Cette fonction esttorch.gatherl'opération inverse de.


Pytorch torch 参考手册Manuel de référence PyTorch torch

Autres extensions