Добавить несколько тензоров на месте в PyTorch

result = torch.sum(torch.stack([x, y, ...]), dim=0)

Без стека:

from functools import reduce

result = reduce(torch.add, [x, y, ...])

ИЗМЕНИТЬ

Как отметил @LudvigH, второй метод не так эффективен с точки зрения памяти, как добавление на место. Так что лучше так:

from functools import reduce

result = reduce(
    torch.Tensor.add_,
    [x, y, ...],
    torch.zeros_like(x)  # optionally set initial element to avoid changing `x`
)

Насколько важно, чтобы операции выполнялись на месте?

Я считаю, что единственный способ выполнить добавление на месте - использовать функцию add_.

Например:

a = torch.randn(5)
b = torch.randn(5)
c = torch.randn(5)
d = torch.randn(5)

a.add_(b).add_(c).add_(d) # in place addition of a+b+c+d

В общем, операции на месте в PyTorch сложны. Они отговаривают его использовать. Я думаю, это связано с тем, что очень легко испортить график вычислений и получить неожиданные результаты. Кроме того, есть много оптимизаций графического процессора, которые в любом случае будут выполнены, и принудительное выполнение операций может в конечном итоге замедлить вашу производительность. Но предполагая, что вы действительно знаете, что делаете, и хотите суммировать множество тензоров с совместимыми формами, я бы использовал следующий шаблон:

import functools
import operator

list_of_tensors = [a, b, c] # some tensors previously defined
functools.reduce(operator.iadd, list_of_tensors)
### now tensor a in the in-place sum of all the tensors

Он основан на шаблоне reduce, что означает сделать это для всех элементов в списке / итерабельном, и operator.iadd, что означает +=. С += есть много предостережений, поскольку он может испортить область видимости и неожиданно вести себя с неизменяемыми переменными, такими как строки. Но в контексте PyTorch он делает то, что мы хотим. Он обращается к add_.

Ниже вы можете увидеть простой тест.

from functools import reduce
from operator import iadd
import torch


def make_tensors():
    return [torch.randn(5, 5) for _ in range(1000)]


def profile_action(label, action):
    print(label)
    list_of_tensors = make_tensors()
    with torch.autograd.profiler.profile(
        profile_memory=True, record_shapes=True
    ) as prof:
        action(list_of_tensors)

    print(prof.key_averages().table(sort_by="self_cpu_memory_usage"))


profile_action("Case A:", lambda tensors: torch.sum(torch.stack(tensors), dim=0))
profile_action("Case B:", lambda tensors: sum(tensors))
profile_action("Case C:", lambda tensors: reduce(torch.add, tensors))
profile_action("Case C:", lambda tensors: reduce(iadd, tensors))

Конечно, результаты различаются между запусками, но на моей машине это копипаст было в некоторой степени типичным. Попробуйте на своем! Вероятно, это немного изменится и с версией pytorch ...

--------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls
--------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
       aten::resize_         0.14%      14.200us         0.14%      14.200us      14.200us      97.66 Kb      97.66 Kb             1
         aten::empty         0.06%       5.800us         0.06%       5.800us       2.900us         100 b         100 b             2
         aten::stack        17.38%       1.751ms        98.71%       9.945ms       9.945ms      97.66 Kb           0 b             1
     aten::unsqueeze        30.55%       3.078ms        78.55%       7.914ms       7.914us           0 b           0 b          1000
    aten::as_strided        48.02%       4.837ms        48.02%       4.837ms       4.833us           0 b           0 b          1001
           aten::cat         0.73%      73.800us         2.78%     280.000us     280.000us      97.66 Kb           0 b             1
          aten::_cat         1.87%     188.900us         2.05%     206.200us     206.200us      97.66 Kb           0 b             1
           aten::sum         1.09%     109.400us         1.29%     130.100us     130.100us         100 b           0 b             1
         aten::fill_         0.17%      16.700us         0.17%      16.700us      16.700us           0 b           0 b             1
            [memory]         0.00%       0.000us         0.00%       0.000us       0.000us     -97.75 Kb     -97.75 Kb             2
--------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 10.075ms

-----------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls
-----------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
              aten::add        99.32%      14.711ms       100.00%      14.812ms      14.812us      97.66 Kb      97.65 Kb          1000
    aten::empty_strided         0.07%      10.400us         0.07%      10.400us      10.400us           4 b           4 b             1
               aten::to         0.37%      54.900us         0.68%     100.400us     100.400us           4 b           0 b             1
            aten::copy_         0.24%      35.100us         0.24%      35.100us      35.100us           0 b           0 b             1
               [memory]         0.00%       0.000us         0.00%       0.000us       0.000us     -97.66 Kb     -97.66 Kb          1002
-----------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 14.812ms

-------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
         Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls
-------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
    aten::add       100.00%      10.968ms       100.00%      10.968ms      10.979us      97.56 Kb      97.56 Kb           999
     [memory]         0.00%       0.000us         0.00%       0.000us       0.000us     -97.56 Kb     -97.56 Kb           999
-------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 10.968ms

--------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
          Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls
--------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
    aten::add_       100.00%       5.184ms       100.00%       5.184ms       5.190us           0 b           0 b           999
--------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 5.184ms

Я выделяю 1000 тензоров, содержащих 25 раз 32-битные числа с плавающей запятой (всего 100 бит на тензор, всего 100 КБ = 97,66 КБ). Разница во времени выполнения и объеме памяти впечатляет.

В случае A torch.sum(torch.stack(list_of_tensors), dim=0) выделяет 100 КБ для стека и 100 КБ для результата, что занимает 10 мс.

Случай B sum занимает 14 мс. Думаю, в основном из-за накладных расходов на Python. Он выделяет 10 КБ для всех промежуточных результатов каждого добавления.

В случае C используется reduce-add, который избавляет от некоторых накладных расходов, увеличивает производительность во время выполнения (11 мс), но при этом выделяет промежуточные результаты. На этот раз он не начинается с 0-инициализации, как это делает sum, поэтому мы выполняем только 999 добавлений вместо 1000 и выделяем на один промежуточный результат меньше. Разница со случаем B незначительна, и в большинстве запусков у них было одинаковое время выполнения.

Случай D - это мой рекомендуемый способ добавления повторяющегося / списка тензоров. Это занимает примерно половину времени и не требует дополнительной памяти. Эффективный. Но вы тратите впустую первый тензор в списке, поскольку выполняете операцию на месте.