深度学习笔记之【einsum】的学习

einsum 是一种强大的矩阵计算工具，在深度学习框架中有实现。
今天在学习源码时看到这样一个式子，理解了很久，看懂了，所以记录一下。
原式如下：

x = torch.einsum('btdc,tw->bwdc', (x, A))

假设我有两个张量，维度分别是

x = torch.randint(1, 3, size=(1, 2, 3, 3))
A = torch.randint(1, 3, size=(2, 3))
print(x)
print(A)

tensor([[[[1, 1, 1],
          [2, 2, 1],
          [2, 2, 2]],

         [[1, 1, 1],
          [1, 1, 2],
          [1, 1, 2]]]])
tensor([[2, 1, 1],
        [1, 1, 2]])

通过einsum定义好规则后计算得到结果：

tensor([[[[3, 3, 3],
          [5, 5, 4],
          [5, 5, 6]],

         [[2, 2, 2],
          [3, 3, 3],
          [3, 3, 4]],

         [[3, 3, 3],
          [4, 4, 5],
          [4, 4, 6]]]])

分析：
btdc,tw->bwdc
$$\begin{gathered} x:(b,t,d,c) \\ A:(tw) \\ \to (b,w,d,c) \\ 可以看到b，c两个维度被保留了下来，而w只出现在A和输出中， \\ 所以w是从A中得到的维度， \\ t同时出现在x和A中而在输出中不存在，说明在这个维度上进行了求和 \end{gathered}$$
从具体的例子中看运算规则：
本例中 t = 2，从 x 中 t 的维度上看就是两个 3 x 3 的矩阵，而从 A 中 t 的维度上看 是一个 2 x 1 的矩阵，举例说明：
从 x 中 t 的维度看到的两个矩阵
$$\left(\begin{array}{ccc}1& 2& 2\\ 2& 2& 2\\ 2& 2& 1\end{array}\right)和\left(\begin{array}{ccc}1& 1& 2\\ 1& 1& 1\\ 2& 2& 2\end{array}\right)$$
从 A 中 t 的维度看到的矩阵
$$\left(\begin{array}{c}2\\ 1\end{array}\right)和\left(\begin{array}{c}2\\ 1\end{array}\right)和\left(\begin{array}{c}2\\ 1\end{array}\right)$$
举其中一个例子：

剩下的两个计算方法相同。
初学者的理解，莫见笑