Home > Back-end >  Torch - How to calculate average of tensors with the same indexes
Torch - How to calculate average of tensors with the same indexes

Time:10-13

Suppose having two matrices: X(m, n) and index matrix I(m, 1). Every item in index matrix I_k represents the index of the kth element X_k in X.

And suppose the index is in the range of [0, 1, 2, ..., j-1]

I would like to calculate the average of tensors in X with the same index i and return a result matrix R(j, n).

For example,

X = [[1, 1, 1],
     [2, 2, 2],
     [3, 3, 3]]

I = [0, 0, 1]

The result matrix should be:

R = [[torch.mean(([1, 1, 1], [2, 2, 2]))],
     [torch.mean(([3, 3, 3]))]

which equals to:

R = [[1.5, 1.5, 1.5],
     [3, 3, 3]]

My current solution is to traverse through m, stack the tensors with the same index and perform torch.mean.

Is there a way avoiding traversing through m? It seems not elegant and rather time-consuming.

CodePudding user response:

ret = torch.empty_like(X)
ret.scatter_reduce_(0, I.unsqueeze(-1).expand_as(X), X, "mean", include_self=False)

should do what you want. Now, note that this is a fairly new method so it may not be particularly performant. If you bump into an issue with this method, you may be better off running scatter_add_ on the tensor X and a tensor of ones and then divide.

If you want to also have a smaller tensor as output, you may want to figure out how many indices and with that infer the size of out.

CodePudding user response:

Cliquer ;

Info)https://atcold.github.io/pytorch-Deep-Learning/fr/week05/05-3/

Astuce : on peut utiliser le question mark dans IPython pour avoir accès aux documents de fonctions et nous arrivons aux mêmes résultats

Fast learn

  • Related