【不写for循环】玩玩行列

利用numpy的并行操作可以比纯用Python的list快很多，不仅如此，代码往往精简得多。

So, 这篇来讲讲进阶的广播和花哨索引操作，少写几个for循环（）。

目录

一个二维的例题

一个三维的例题

解法一

解法二

更难的三维例题

解法一

解法二

独热编码

写法一

写法二

一个二维的例题

从一个简单的问题开始，现在有一个向量：

弄出这个东西出来：

这个很简单：

x = [1, 2, 3]
res = []
for i in x:for j in x:res.append(i + j)
res = torch.tensor(res).reshape(3, 3)
res# output
tensor([[2, 3, 4],[3, 4, 5],[4, 5, 6]])

以上内容是开玩笑的，下面开始认真起来。自然，这个很明显是个广播的送分题：

x = torch.tensor([1, 2, 3])
x + x.reshape(-1, 1)#output
tensor([[2, 3, 4],[3, 4, 5],[4, 5, 6]])

一个三维的例题

现在有一个向量：

弄出这个东西出来：

不过现在a b c都是长度为4的向量。x是一个(3, 4)的矩阵（还是说明一下，这里不表示分块矩阵），目标是一个(3, 3, 4)的张量。

解法一

这个在GAT里面里面很常见（把“+”换成“concat”就是图卷积注意力的核心步骤之一）。当时看了一大圈的zhihu和CSDN，都是这么写的：

先x.repeat(1,3)，横着重复，维度是(3, 3*4):

（||表示两个向量拼接）

然后x.reshape(3*3, -1), 维度变成(3*3, 4):

另一个竖着重复，x.repeat(3,1)，维度是(3*3, 4):

然后相加reshape即可。

x = torch.tensor([[1, 1, 1, 1],[2, 2, 2, 2],[3, 3, 3, 3]])
(x.repeat(1, 3).reshape(3*3, 4) + x.repeat(3, 1)).reshape(3, 3, 4)#output
tensor([[[2, 2, 2, 2],[3, 3, 3, 3],[4, 4, 4, 4]],[[3, 3, 3, 3],[4, 4, 4, 4],[5, 5, 5, 5]],[[4, 4, 4, 4],[5, 5, 5, 5],[6, 6, 6, 6]]])

解法二

后来想了一下其实可以直接广播。

解法一虽然复杂一点，但是把题目里面的“+”改成“||”就只能用解法一了。

x = torch.tensor([[1, 1, 1, 1],[2, 2, 2, 2],[3, 3, 3, 3]])
x.unsqueeze(0) + x.unsqueeze(1)  # (1, 3, 4) + (3, 1, 4)#output
tensor([[[2, 2, 2, 2],[3, 3, 3, 3],[4, 4, 4, 4]],[[3, 3, 3, 3],[4, 4, 4, 4],[5, 5, 5, 5]],[[4, 4, 4, 4],[5, 5, 5, 5],[6, 6, 6, 6]]])

更难的三维例题

现在有一个向量：

弄出这个东西出来：

不过现在a b c都是长度为4的向量。x是一个(3, 4)的矩阵（还是说明一下，这里不表示分块矩阵），两两做点积，目标是一个(3, 3)的张量。

先来一个错误示例：

x = torch.tensor([[1, 1, 1, 1],[2, 2, 2, 2],[3, 3, 3, 3]])
np.dot(x.unsqueeze(0), x.unsqueeze(1)), torch.dot(x.unsqueeze(0), x.unsqueeze(1))

两种做法都是错的，torch.dot只支持1D的向量。np,dot处理高维度的张量的逻辑很不同，这里可以

查阅资料，不细说了。

解法一

可以用numpy里面最玄学的函数之一——np.meshgrid

先看看这个函数是干嘛的：

x = torch.tensor([[1, 1, 1, 1],[2, 2, 2, 2],[3, 3, 3, 3]])
i, j = np.meshgrid(np.arange(x.shape[0]), np.arange(x.shape[1]), indexing='ij')
i, j# output
array([[0, 0, 0, 0],[1, 1, 1, 1],[2, 2, 2, 2]]array([[0, 1, 2, 3],[0, 1, 2, 3],[0, 1, 2, 3]])

np.arange(x.shape[0]) : array([0, 1, 2])

np.arange(x.shape[0]) : array([0, 1, 2, 3])

然后这个函数让前者往右重复，让后者往下重复，得到两个矩阵。然后细心看花哨索引和广播就知道：

x == x[i, j]  !!!

了解这个函数干嘛后，那下面我们进入正题。

x = torch.tensor([[1, 1, 1, 1],[2, 2, 2, 2],[3, 3, 3, 3]])
x1, x2 = x.unsqueeze(0), x.unsqueeze(1)  # (1, 3, 4) (3, 1, 4)
x1, x2 = torch.broadcast_tensors(x1, x2)  # (3, 3, 4) (3, 3, 4) 手动广播
i, j = np.meshgrid(np.arange(3), np.arange(3), indexing='ij')
torch.sum(x1[i, j, :] * x2[i, j, :], dim=-1)#output
tensor([[ 4,  8, 12],[ 8, 16, 24],[12, 24, 36]])

用花哨索引固定前两个维度不动，在第三个维度上相乘求和（就是点积）。搞定。

解法二

x = torch.tensor([[1, 1, 1, 1],[2, 2, 2, 2],[3, 3, 3, 3]])
x1, x2 = x.unsqueeze(0), x.unsqueeze(1)  # (1, 3, 4) (3, 1, 4)
torch.einsum('ijk,ijk->ij', x1, x2)  # 这个函数支持广播#output
tensor([[ 4,  8, 12],[ 8, 16, 24],[12, 24, 36]])

np.einsum的全称是Einstein summation convention，即爱因斯坦求和约定。这个约定允许我们通过一个简洁的字符串表达式来指定复杂的数组运算，包括点积、矩阵乘法、张量收缩等。

这里是一个简单的运用。

独热编码

原来利用广播可以写独热编码。

写法一

一般独热编码可以这么写

a = np.array([1, 2, 1, 0])
category = len(np.unique(a))
eye = np.eye(category)
eye, eye[a]#output
array([[1., 0., 0.],[0., 1., 0.],[0., 0., 1.]]array([[0., 1., 0.],[0., 0., 1.],[0., 1., 0.],[1., 0., 0.]]

eye是一个单位矩阵，a构成了一个花哨索引，每次取eye的一行，然后取4次。

十分简洁。缺点是a的每个值必须在[0,category-1]中。

写法二

a = np.array(['a', 'b', 'c', 'd', 'e', 'f'])
b = np.array(['d', 'e', 'f'])
b = b.reshape(-1, 1)  # (3, 1)
(a == b).astype(int)# output
array([[0, 0, 0, 1, 0, 0],[0, 0, 0, 0, 1, 0],[0, 0, 0, 0, 0, 1]]

支持各种类型的数据，而且还能应对b的某个元素不在a中的尴尬情况（此时一排都是0，因为一排都是不等于）。

有一个缺点是，在第四行时，Pycharm不知道这是一个a==b是一个布尔数组，在"astype"会画一个黄色，看着闹心（狗头）。