I can add two tensors x and y inplace like this
x = x.add(y)
Is there a way of doing the same with three or more tensors given all tensors have same dimensions?
Add multiple tensors inplace in PyTorch
Viewed 4519
3 Answers
result = torch.sum(torch.stack([x, y, ...]), dim=0)
Without stack:
from functools import reduce
result = reduce(torch.add, [x, y, ...])
EDIT
As @LudvigH pointed out, the second method is not as memory-efficient, as inplace addition. So it's better like this:
from functools import reduce
result = reduce(
torch.Tensor.add_,
[x, y, ...],
torch.zeros_like(x) # optionally set initial element to avoid changing `x`
)
How important is it that the operations occur in place?
I believe the only way to do addition in place is with the add_ function.
For example:
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
In general, in-place operations in PyTorch is tricky. They discourage using it. I think it stems from the fact that it is easy to mess up, ruinig the computation graph and giving you unexpected results. Also, there are many GPU optimizations that will be done anyway, and forcing in place operations can slow down your performance in the end. But assuming that your really know what you are doing, and you want to sum a lot of tensors with compatible shapes I would use the following pattern:
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
It builds on the pattern of reduce, which means "do this to all elements in the list/iterable", and operator.iadd which means +=. There are many caveats with +=, since it may mess up with scoping and behaves unexpectedly with immutable variables such as strings. But in the context of PyTorch, it does what we want. It makes calls to add_.
Below, you can see a simple benchmark.
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))
The results vary between runs of course, bit this copy-paste was somewhat representative on my machine. Try it on yours! It probably changes a bit with pytorch version as well...
-------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
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
I allocate 1000 tensors holding 25 times 32 bit floats (total 100b per tensor, 100kb=97.66Kb in total). The difference in runtime and memory footprint is quite stunning.
Case A, torch.sum(torch.stack(list_of_tensors), dim=0) allocates 100kb for the stack, and 100b for the result, taking 10 ms.
Case B sum takes 14ms. Mostly because of python overhead I guess. It allocates 10kb for all intermediary results of each addition.
Case C uses reduce-add which gets rid of some overhead, gaining in runtime performance (11ms) but still allocating intermediary results. This time, it does not start with a 0-initialization, which sum does, so we only do 999 additions instead of 1000, and allocate one intermediate result less. The difference with Case B is minute, and in most runs they had the same runtime.
Case D is my recommended way for in place addition of an iterable/list of tensors. It takes roughly half the time and allocates no extra memory. Efficient. But you waste the first tensor in the list, since you perform the operation in place.
Related