Pytorch bincount with gradient

Question

I am trying to get gradient from sum of some indexes of an array using bincount. However, pytorch does not implement the gradient. This can be implemented by a loop and torch.sum but it is too slow. Is it possible to do this efficiently in pytorch (maybe einsum or index_add)? Of course, we can loop over indexes and add one by one, however that would increase the computational graph size significantly and is very low performance.

import torch
from torch import autograd
import numpy as np
tt = lambda x, grad=True: torch.tensor(x, requires_grad=grad)    
inds = tt([1, 5, 7, 1], False).long()
y = tt(np.arange(4) + 0.1).float()
sum_y_section = torch.bincount(inds, y * y, minlength=8)
#sum_y_section = torch.sum(y * y)
grad = autograd.grad(sum_y_section, y, create_graph=True, allow_unused=False)
print("sum_y_section", sum_y_section)
print("grad", grad)

Have you tried implementing it with a slow loop implementation as you described? — Ivan, May 02 '22 at 15:23

Roy · Accepted Answer · 2022-05-24T23:20:21.563

3

We can use a new feature in Pytorch V1.11 called scatter_reduce.

bincount = lambda inds, arr: torch.scatter_reduce(arr, 0, inds, reduce="sum")

edited May 24 '22 at 23:20

answered May 03 '22 at 03:57

Roy

65
2
15
40

4

it would be nice if you could add a code snippet of the solution with `scatter_reduce` - for future readers of this answer... – Shai May 03 '22 at 05:37
1

Sure, I will add. – Roy May 03 '22 at 07:23

score 1 · Answer 2 · answered May 06 '22 at 15:14

1

I’d try to use a hook to manipulate the gradient in a custom way

answered May 06 '22 at 15:14

emely_pi

517
5
23

score 1 · Answer 3 · answered Mar 20 '23 at 19:21

torch.scatter_reduce has a src positional argument in Pytorch 1.13. This simple example demonstrates equivalence between bincount and scatter_reduce:

num_bins = 8
bins = torch.zeros(num_bins)

#generate indices and weights
num_indices = 100
indices = torch.randint(num_bins, size=(num_indices,))
weights = torch.rand(num_indices)

# Counting Indices

# with torch.bincount
counts1 = torch.bincount(indices, minlength=num_bins)

# with torch.scatter_reduce
counts2 = bins.scatter_reduce(0, indices, torch.ones(num_indices), reduce = 'sum')
print(counts1)
print(counts2)

# Binning Weights

# with torch.bincount
binned_wts1 =  indices.bincount(weights, minlength=num_bins)

# with torch.scatter_reduce
binned_wts2 = bins.scatter_reduce(0, indices, weights, reduce='sum')

print(binned_wts1)
print(binned_wts2)

Pytorch bincount with gradient

3 Answers3