你不知道的 Numpy—einsum

这是我参与2022首次更文挑战的第9天，活动详情查看：2022首次更文挑战

今天要说的是 einsum ，可能是接触过 Numpy 也很少会用到这个。einsum全称Einstein summation convention。在速度和内存效率方面往往超过了熟悉的数组函数。

在正式开始介绍 einsum 之前，我们先演示一些例子，例如定义一个函数rnd，利用这个函数可以生成多维一个矩阵，矩阵的元素是随机数，调用 rnd 只需要指定我们想要矩阵的形状(shape)，也就是具体指出矩阵维数以及每一个维度的长度，然后这个函数就可以随机生成你需要的矩阵。

def rnd(*args):
  return np.random.random(args)

Numpy 提供了 matmul 方法用于对两个矩阵进行相乘，如果要对两个矩阵相乘，需要两个矩阵满足一定条件，也就是第一个矩阵列的维度和第二个矩阵行的维度相同。这里再多说一句，其应用场景可以是对一个向量或者矩阵进行线性变换。或者在隐含层 $y = WX^T$

np.matmul(rnd(4,5),rnd(5,6)).shape
#(4, 6)

通过 rnd 函数得到两个矩阵 A 和 B 然后将 B 进行转置，让 B 的形状满足可以与 A 相乘。然后在将 A 和 B 相乘。

A = rnd(5,32)
B = rnd(10,32)
np.matmul(A,B.T).shape
#(5, 10)

np.matmul(rnd(10,4,5),rnd(10,5,6)).shape
#(10, 4, 6)

对于matmul方法来说，只对最后两个维度进行矩阵相乘，其他维度保持不变，上面代码可以理解为一个样本批量为 10 的小批量样本输入到，接下来进一步扩展到更高维度上两个矩阵相乘，

np.matmul(rnd(30,20,10,4,5),rnd(30,20,10,5,6)).shape

下面代码也是比较实用，在 transformer 我们 query 和 keys 进行，也就是用一个 token 的 query 分别与整句或者图像像素分别相乘来计算他们之间相似度，或者叫相关性。可以通过将两个向量相乘来计算其相关性，因为有 10 个这样向量所以将分别和向量相乘计算转换为一个矩阵和向量相乘

E = 32
query_embedding = rnd(E)
keys_embeddings = rnd(10,E)
v1 = np.matmul(query_embedding,keys_embeddings.T)
v2 = np.matmul(keys_embeddings,query_embedding)
v1.shape, v2.shape, np.all(np.isclose(v1,v2))
#((10,), (10,), True)

然后我们添加一个维度表示每一个批量样本的数量，这里 N = 100

N = 100
E = 32
query_embedding = rnd(N,E)
keys_embeddings = rnd(N,10,E)

我们需要给给 query 添加一个维度，

query_embedding.reshape(N,1,E).shape
# (100, 1, 32)

还需要对 key 进行维度 reshape 这样 query 和 key 才能够相乘

keys_embeddings.transpose(0,2,1).shape
#(100, 32, 10)

np.squeeze(np.matmul(query_embedding.reshape(N,1,E),keys_embeddings.transpose(0,2,1))).shape
# (100, 10)

接下来我们就开始将上面代码用 einsum 一起实现，einsum 这里 ij 分别代表 A 的行和列，而矩阵 B 为行应该和 A 列维度长度一致，所以相乘矩阵得到维度为 ik

A=rnd(4,5)
B=rnd(5,6)

m = np.matmul(A,B)
e = np.einsum('ij,jk->ik',A,B)
m.shape,e.shape,np.all(np.isclose(m,e))
# ((4, 6), (4, 6), True)

A=rnd(5,32)
B=rnd(10,32)

m = np.matmul(A,B.T)
e = np.einsum('ij,kj->ik',A,B)
m.shape,e.shape,np.all(np.isclose(m,e))
# ((5, 10), (5, 10), True)