带梯度的 Pytorch bincount

Question

我正在尝试使用 bincount 从数组的某些索引的总和中获取梯度。但是，pytorch 没有实现渐变。这可以通过循环和 torch.sum 来实现，但它太慢了。是否可以在 pytorch 中有效地执行此操作（可能是 einsum 或 index_add）？当然，我们可以循环遍历索引并逐个添加，但是这会显着增加计算图的大小并且性能非常低。

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)

Answer 1

我们可以使用 Pytorch V1.11 中称为 scatter_reduce 的新功能。

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

Answer 2

我会尝试使用钩子以自定义方式操纵渐变

带梯度的 Pytorch bincount

Pytorch bincount with gradient

pytorch

autograd