问题描述
这是有关torch.einsum
在GPU中内部工作的查询。我知道如何使用einsum
。它执行所有可能的矩阵乘法,然后只选择相关的乘法,还是仅执行所需的计算?
例如,考虑形状为a
的两个张量b
和(N,P)
,我希望找到形状为{的每个张量ni
的点积{1}}。
使用einsum,代码为:
(1,P)
torch.einsum('ij,ij->i',a,b)
现在,第二个代码应该比第一个执行更多的计算(例如,如果torch.diag(a @ b.t())
= N
,则它执行的计算要多2000
倍)。但是,当我尝试对这两个操作进行计时时,它们花费的时间大致相同,这引出了问题。 2000
是否执行所有组合(如第二个代码),并选择相关值?
要测试的示例代码:
einsum
解决方法
这可能与GPU可以并行计算a @ b.t()
的事实有关。这意味着GPU实际上不必等待每个行-列乘法计算完成就可以计算然后进行下一个乘法。
如果检查CPU,则会发现torch.diag(a @ b.t())
和torch.einsum('ij,ij->i',a,b)
对于大型a
和b
的速度要慢得多。
我不能代表torch
,但是几年前与np.einsum
合作过。然后,它基于索引字符串构造了一个自定义迭代器,仅执行必要的计算。从那时起,它已经以各种方式进行了重新设计,并在可能的情况下将问题显然转换为@
,从而利用了BLAS(etc)库调用的优势。
In [147]: a = np.arange(12).reshape(3,4)
In [148]: b = a
In [149]: np.einsum('ij,b)
Out[149]: array([ 14,126,366])
我不确定在这种情况下使用哪种方法。使用“ j”求和,也可以使用:
In [150]: (a*b).sum(axis=1)
Out[150]: array([ 14,366])
您注意到,最简单的dot
创建了一个更大的数组,我们可以从中拉出对角线:
In [151]: (a@b.T).shape
Out[151]: (3,3)
但这不是使用@
的正确方法。 @
通过提供有效的“批处理”来扩展np.dot
。因此,i
维度是批次一,j
是dot
维度。
In [152]: a[:,None,:]@b[:,:,None]
Out[152]:
array([[[ 14]],[[126]],[[366]]])
In [156]: (a[:,None])[:,0]
Out[156]: array([ 14,366])
换句话说,它使用(3,1,4)与(3,4,1)来产生(3,1),对尺寸为4的共享乘积求和。
一些采样时间:
In [162]: timeit np.einsum('ij,b)
7.07 µs ± 89.2 ns per loop (mean ± std. dev. of 7 runs,100000 loops each)
In [163]: timeit (a*b).sum(axis=1)
9.89 µs ± 122 ns per loop (mean ± std. dev. of 7 runs,100000 loops each)
In [164]: timeit np.diag(a@b.T)
10.6 µs ± 31.4 ns per loop (mean ± std. dev. of 7 runs,100000 loops each)
In [165]: timeit (a[:,0]
5.18 µs ± 197 ns per loop (mean ± std. dev. of 7 runs,100000 loops each)