python - 如何向量化 scatter-matmul 运算

Question

我有许多带有形状的矩阵 w1、w2、w3...wn (k*n1 、k*n2、k*n3...k*nn) 和 x1、x2、x3...xn 具有形状(n1*m、n2*m、n3*m...nn*m >).
我想分别获取 w1@x1、w2@x2、w3@x3 ...。

生成的矩阵是多个 k*m 矩阵，可以连接成形状为 (k*n)*m 的大矩阵。

将它们一一相乘会很慢。如何向量化这个操作？

注意:输入可以是 k*(n1+n2+n3+...+nn) 矩阵和 (n1+n2+n3+...+nn)* m 矩阵，我们可以使用批量索引来指示这些子矩阵。

此操作与pytorch_scatter中实现的分散操作相关。 ，所以我将其称为“scatter_matmul”。

Answer 1

您可以通过创建形状为 n*kx(n1+..+nn) 的大型 block 对角矩阵 W 来矢量化您的操作> 其中 w_i 矩阵是对角线上的 block 。然后，您可以将所有 x 矩阵垂直堆叠成形状为 (n1+..+nn)xm 的 X 矩阵>。将 block 对角线 W 与所有 x 矩阵的垂直堆栈相乘，X:

Y = W @ X

Y 形状为 (k*n)xm 的结果，这正是您正在寻找的串联大矩阵。

如果 block 对角矩阵W的形状太大而无法放入内存，您可以考虑将W sparse并使用 torch.sparse.mm 计算乘积.