Python 向量化不同维度的 2d 和 3d 数组串联

Question

我正在尝试在图神经网络中实现消息传递。在每个图中，都有边和节点，节点到边的更新实现如下：其中方括号表示拼接操作，下标为索引，上标为时间索引。

所以我试图连接 3 个维度矩阵：AxN、AxBxM 和 BxN。并且由此产生的串联的维度是：AxBx(2N+M)。因此，结果矩阵的每个 (i,j) 都是第一个矩阵的第 i 行、第三个矩阵的第 j 行和第二个矩阵的第 (i,j) 个元素的串联。我设法在双 for 循环中实现了它，如下所示：

edge_in = torch.zeros(a, b, m +  2 * n)
        edge_in = edge_in.cuda()
        for i in range(a):
          for j in range(b):
              edge_in[i,j] = torch.cat((nodes_a_embeds[i], edge_embeds[i,j], nodes_b_embeds[j]))

但是，这非常慢。这是可以向量化的吗？我试图想出一个解决方案，然后我在网上寻找解决方案，但无法对其进行矢量化。谢谢

编辑：请求的数值示例：

第一个矩阵：5x3 第二个矩阵：5x4x2 第三个矩阵：4x3

那么输出应该是 5x4x8。我们称我们的输出矩阵为 R.

然后 R(1,2) = concatenate(First(1),Second(1,2),Third(2)).

Answer 1

这是您代码的正确实现吗？

import numpy as np

A = 2
B = 3
M = 4
N = 5

first = np.arange(A*N).reshape((A, N))
first = np.tile(first[:, np.newaxis, :], (1, B, 1))

second = np.arange(A*B*M).reshape((A, B, M))

third = np.arange(B*N).reshape((B, N))
third = np.tile(third[np.newaxis, :, :], (A, 1, 1))

result = np.concatenate((first, second, third), axis=2)

输出：

array([[[ 0,  1,  2,  3,  4,  0,  1,  2,  3,  0,  1,  2,  3,  4],
        [ 0,  1,  2,  3,  4,  4,  5,  6,  7,  5,  6,  7,  8,  9],
        [ 0,  1,  2,  3,  4,  8,  9, 10, 11, 10, 11, 12, 13, 14]],

       [[ 5,  6,  7,  8,  9, 12, 13, 14, 15,  0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9, 16, 17, 18, 19,  5,  6,  7,  8,  9],
        [ 5,  6,  7,  8,  9, 20, 21, 22, 23, 10, 11, 12, 13, 14]]])

Python 向量化不同维度的 2d 和 3d 数组串联

Python vectorize 2d and 3d array concatenation with different dimensions

python

concatenation

matrix

vectorization