QBoard » Artificial Intelligence & ML » AI and ML - PyTorch » Torch sum a tensor along an axis
User Dashboard

Torch sum a tensor along an axis

Back To Topics
Tags : None
  • 84 
    ipdb> outputs.size()
torch.Size([10, 100])
ipdb> print sum(outputs,0).size(),sum(outputs,1).size(),sum(outputs,2).size()
(100L,) (100L,) (100L,)

    How do I sum over the columns instead?

      September 23, 2021 11:10 PM IST
    0
  • 461
    If you have tensor my_tensor, and you wish to sum across the second array dimension (that is, the one with index 1, which is the column-dimension, if the tensor is 2-dimensional, as yours is), use torch.sum(my_tensor,1) or equivalently my_tensor.sum(1) see documentation here.

    One thing that is not mentioned explicitly in the documentation is: you can sum across the last array-dimension by using -1 (or the second-to last dimension, with -2, etc.)

    So, in your example, you could use: outputs.sum(1) or torch.sum(outputs,1), or, equivalently, outputs.sum(-1) or torch.sum(outputs,-1). All of these would give the same result, an output tensor of size torch.Size([10]), with each entry being the sum over the all rows in a given column of the tensor outputs.

    To illustrate with a 3-dimensional tensor:

    In [1]: my_tensor = torch.arange(24).view(2, 3, 4) 
Out[1]: 
tensor([[[ 0,  1,  2,  3],
         [ 4,  5,  6,  7],
         [ 8,  9, 10, 11]],

        [[12, 13, 14, 15],
         [16, 17, 18, 19],
         [20, 21, 22, 23]]])

In [2]: my_tensor.sum(2)
Out[2]:
tensor([[ 6, 22, 38],
        [54, 70, 86]])

In [3]: my_tensor.sum(-1)
Out[3]:
tensor([[ 6, 22, 38],
        [54, 70, 86]])​
      September 24, 2021 1:36 PM IST
    0
  • 63 1
    Alternatively, you can use tensor.sum(axis) where axis indicates 0 and 1 for summing over rows and columns respectively, for a 2D tensor.

    In [210]: X
Out[210]: 
tensor([[  1,  -3,   0,  10],
        [  9,   3,   2,  10],
        [  0,   3, -12,  32]])

In [211]: X.sum(1)
Out[211]: tensor([ 8, 24, 23])

In [212]: X.sum(0)
Out[212]: tensor([ 10,   3, -10,  52])​

    As, we can see from the above outputs, in both cases, the output is a 1D tensor. If you, on the other hand, wish to retain the dimension of the original tensor in the output as well, then you've set the boolean kwarg keepdim to True as in:

    In [217]: X.sum(0, keepdim=True)
Out[217]: tensor([[ 10,   3, -10,  52]])

In [218]: X.sum(1, keepdim=True)
Out[218]: 
tensor([[ 8],
        [24],
        [23]])
      September 25, 2021 12:05 AM IST
    0
  • 317
    The simplest and best solution is to use torch.sum().

    To sum all elements of a tensor:

    torch.sum(outputs) # gives back a scalar
​

    To sum over all rows (i.e. for each column):


    torch.sum(outputs, dim=0) # size = [1, ncol]
​


    To sum over all columns (i.e. for each row):

    torch.sum(outputs, dim=1) # size = [nrow, 1]
      September 28, 2021 1:44 PM IST
    0