python - 基于索引的结构张量

Question

我正在尝试根据相应的索引张量连接多个张量。作为一个玩具示例，我有一个张量 A = [[0,1,2],[2,3,4]] 和相应的索引向量 B = [0,1] 和另一个张量 C = [[5,6,7],[8,9,10],[11,12,13]] 以及相应的索引向量 D = [0,1,2]。我寻求返回一个张量，使得 E = [[0,1,2],[5,6,7],[2,3,4],[8,9,10],[11,12 ,13]] 其中具有相同索引的值被分组在一起。我正在尝试有效地执行此操作，因为它包含在我的神经网络模型中，并且在反向传播期间可能会遇到性能问题。

我能够通过构建字典和 for 循环来实现这一点:

E = defaultdict(list)
for vector, indices in zip([A,C],[B,D]):
    for vector, i in enumerate(vector):
        E[indices[i]].append(vector)
E = list(E.values())

然而，这种方法会显着减慢模型训练速度。有没有办法更有效地做到这一点？ Pytorch 的各种 index_ 操作似乎没有解决这个问题。谢谢!

Answer 1

我假设 A 和 C 张量具有不同的行数，但列数相同。

A = torch.LongTensor([[0,1,2],[2,3,4]]) # torch.Size([2, 3])
C = torch.LongTensor([[5,6,7],[8,9,10],[11,12,13]]) # torch.Size([3, 3])

assert A.dim() == C.dim() == 2
assert A.size(1) == C.size(1)

shape_E = (A.size(0) + C.size(0), A.size(1))
E = torch.zeros(*shape_E).long()

# Two possible cases
if A.size(0) > C.size(0): # C has fewer rows
    n_rows = C.size(0)
    E[np.arange(0, 2*n_rows, 2), :] = A[:n_rows, :]
    E[np.arange(1, 2*n_rows, 2), :] = C[:n_rows, :]
    E[2*n_rows:, :] = A[n_rows:, :]
else: # A has fewer rows
    n_rows = A.size(0)
    E[np.arange(0, 2*n_rows, 2), :] = A[:n_rows, :]
    E[np.arange(1, 2*n_rows, 2), :] = C[:n_rows, :]
    E[2*n_rows:, :] = C[n_rows:, :]

print(E)

输出

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