我有 2 个张量。第一个张量是 1D(例如 3 个值的张量)。第二个张量是二维的,第一个暗淡作为一对多关系中第一个张量的 ID(例如形状为 6, 2 的张量)
# e.g. simple example of dot product
import torch
a = torch.tensor([2, 4, 3])
b = torch.tensor([[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 = [(2*2)+(2*3)+(2*1),(4*4),(3*3)+(3*1)]
output = [12, 16, 12]
目前我所拥有的是找到 b 中每个 id 的大小(例如 [3,1,2]),然后使用 torch.split 将它们分组到张量列表中,并在组中运行 for 循环。对于小张量来说还好,但是当张量的大小为数百万个、具有数万个任意大小的组时,它变得非常慢。
还有更好的解决方案吗?
最佳答案
您可以使用numpy.bincount
或torch.bincount
按键对b
的元素求和:
import numpy as np
a = np.array([2,4,3])
b = np.array([[0,2], [0,3], [0,1], [1,4], [2,3], [2,1]])
print( np.bincount(b[:,0], b[:,1]) )
# [6. 4. 4.]
print( a * np.bincount(b[:,0], b[:,1]) )
# [12. 16. 12.]
import torch
a = torch.tensor([2,4,3])
b = torch.tensor([[0,2], [0,3], [0,1], [1,4], [2,3], [2,1]])
torch.bincount(b[:,0], b[:,1])
# tensor([6., 4., 4.], dtype=torch.float64)
a * torch.bincount(b[:,0], b[:,1])
# tensor([12., 16., 12.], dtype=torch.float64)
引用文献:
关于algorithm - pytorch:对两个不同大小的张量执行运算的有效方法,其中一个张量具有一对多关系,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/75101543/