Функция torch.scatter_reduce в PyTorch

Pytorch torch 参考手册Справочное руководство PyTorch torch

torch.scatter_reduceЭто функция в PyTorch, используемая для агрегирования значений исходного тензора в указанные позиции заданным способом. Она поддерживает несколько методов агрегирования, таких как sum, product, max, min и т. д.

Определение функции

torch.scatter_reduce(input, dim, index, src, reduce='sum', *, include_self=True)

Parameters:

  • input(Tensor): Входной тензор.
  • dim(int): Размерность, вдоль которой выполняется агрегирование.
  • index(Tensor): Тензор индексов, указывающий, куда агрегировать значения src в input.
  • src(Tensor): Исходный тензор, значения для агрегации.
  • reduce(str): Метод агрегирования. Варианты: 'sum', 'prod', 'mean', 'amax', 'amin', 'multiply'. По умолчанию — 'sum'.
  • include_self(bool, необязательно): Включать ли в агрегирование исходное значение в самой позиции индекса. По умолчанию — True.

Возвращаемое значение:

  • torch.Tensor: Возвращает агрегированный тензор.

Примеры использования

Пример

import torch

# Создать входной тензор
input = torch.ones(3, 5)

# Создать index и 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]])

# Использовать агрегирование sum
output = torch.scatter_reduce(input, dim=0, index=index, src=src, reduce='sum')

print("Вход:")
print(input)
print("nМетод агрегирования: sum")
print("Результат:")
print(output)

Результат вывода:

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

聚合方式: sum
结果:
tensor([[4., 2., 3., 4., 4.],
        [2., 2., 2., 2., 2.],
        [3., 3., 2., 3., 3.]])

Пример

import torch

# Проверить различные методы агрегации
input = torch.zeros(3, 3)
index = torch.tensor([[0, 0, 0],
                      [1, 1, 1],
                      [2, 2, 2]])
src = torch.tensor([[2, 3, 4],
                    [5, 6, 7],
                    [8, 9, 10]])

# Использовать prod (произведение)
output_prod = torch.scatter_reduce(input, 0, index, src, reduce='prod')
print("агрегация prod:")
print(output_prod)

# Использовать amax (максимум)
output_max = torch.scatter_reduce(input, 0, index, src, reduce='amax')
print("агрегация namax:")
print(output_max)

# Использовать amin (минимум)
output_min = torch.scatter_reduce(input, 0, index, src, reduce='amin')
print("агрегация namin:")
print(output_min)

Результат вывода:

prod 聚合:
tensor([[ 2.,  3.,  4.],
        [ 5.,  6.,  7.],
        [ 8.,  9., 10.]])

amax 聚合:
tensor([[ 2.,  3.,  4.],
        [ 5.,  6.,  7.],
        [ 8.,  9., 10.]])

amin 聚合:
tensor([[ 2.,  3.,  4.],
        [ 5.,  6.,  7.],
        [ 8.,  9., 10.]])

Пример

import torch

# Использование параметра include_self
input = torch.tensor([1.0, 2.0, 3.0])
index = torch.tensor([0, 0, 0])
src = torch.tensor([10.0, 20.0, 30.0])

# Включить собственное значение (по умолчанию)
output1 = torch.scatter_reduce(input, 0, index, src, reduce='sum', include_self=True)
print("include_self=True:", output1)

# Исключить собственное значение
output2 = torch.scatter_reduce(input, 0, index, src, reduce='sum', include_self=False)
print("include_self=False:", output2)

Результат вывода:

include_self=True: tensor([66.,  2.,  3.])
include_self=False: tensor([60.,  2.,  3.])

Примечание:torch.scatter_reduceОн не изменяет исходный входной тензор, а возвращает новый тензор. Несколько индексов могут указывать на одну и ту же позицию, и значения будут агрегированы в соответствии с указанным методом агрегации.include_self=FalseЭто полезно в некоторых сценариях, например, чтобы избежать повторного вычисления собственных признаков узла в графовых нейронных сетях.

Pytorch torch 参考手册Справочное руководство PyTorch torch

Другие расширения