我有两个Tensor,第一个Tensor是1D(例如,3个值的Tensor),第二个Tensor是2D,第一个dim作为第一个Tensor的ID,具有一对多的关系(例如,形状为6,2的Tensor)
# e.g. simple example of dot product
a = [2,4,3]
b = [[0,2], [0,3], [0,1], [1,4], [2,3], [2,1]] # 1st column is the index to tensor a, 2nd column is the value
output = [12,16,12]
目前我所做的是找到B中每个id的大小(例如[3,1,2]),然后使用torch.split将它们分组为一个Tensor列表,并在这些组中运行一个for循环。这对于一个小Tensor来说是可以的,但当Tensor的大小以百万计,有数万个任意大小的组时,它就变得非常慢。
有更好的解决办法吗?
1条答案
按热度按时间gywdnpxw1#
可以使用
numpy.bincount
按键对b
的元素求和:参考文献: