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]]的索引相同的值组合在一起。我正在尝试有效地执行此操作，因为它包含在我的NN模型中，并且在向后传播期间可能会遇到性能问题。

我能够通过构造一个字典并像这样循环来实现：

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]])

基于张量的结构张量

1 个答案: