火炬/ np einsum在内部如何工作

如何解决火炬/ np einsum在内部如何工作

这是有关torch.einsum在GPU中内部工作的查询。我知道如何使用einsum。它执行所有可能的矩阵乘法，然后只选择相关的乘法，还是仅执行所需的计算？

例如，考虑形状为a的两个张量b和(N,P)，我希望找到形状为{的每个张量ni的点积{1}}。 使用einsum，代码为：

(1,P)

不使用einsum，获取输出的另一种方法是：

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)