如何在pytorch中有效地重复时间的张量元素变量?

How to efficiently repeat tensor element variable of time in pytorch?

例如,如果我有一个张量 A = [[1,1,1], [2,2,2], [3,3,3]],并且 B = [1,2,3 ].我如何获得 C = [[1,1,1], [2,2,2], [2,2,2], [3,3,3], [3,3,3], [3, 3,3]],然后分批进行?

顺便说一句,我当前的元素明智的解决方案(永远……):

        def get_char_context(valid_embeds, words_lens):
            chars_contexts = []
            for ve, wl in zip(valid_embeds, words_lens):
                for idx, (e, l) in enumerate(zip(ve, wl)):
                    if idx ==0:
                        chars_context = e.view(1,-1).repeat(l, 1)
                    else:
                        chars_context = torch.cat((chars_context, e.view(1,-1).repeat(l, 1)),0)
                chars_contexts.append(chars_context)
            return chars_contexts

我这样做是为了将 bert 词嵌入添加到字符级 seq2seq 任务中...

使用这个:

import torch
# A is your tensor
B = torch.tensor([1, 2, 3])
C = A.repeat_interleave(B, dim = 0)

编辑:

如果 A 是一个单一的二维张量,上面的方法工作正常。要以相同的方式批量重复所有 (2D) 张量,这是一个简单的解决方法:

A = torch.tensor([[[1, 1, 1], [2, 2, 2], [3, 3, 3]], 
    [[1, 2, 3], [4, 5, 6], [2,2,2]]]) # A has 2 tensors each of shape (3, 3)
B = torch.tensor([1, 2, 3]) # Rep. of each row of every tensor in the batch

A1 = A.reshape(1, -1, A.shape[2]).squeeze()
B1 = B.repeat(A.shape[0])
C = A1.repeat_interleave(B1, dim = 0).reshape(A.shape[0], -1, A.shape[2])

C 是:

tensor([[[1, 1, 1],
         [2, 2, 2],
         [2, 2, 2],
         [3, 3, 3],
         [3, 3, 3],
         [3, 3, 3]],

        [[1, 2, 3],
         [4, 5, 6],
         [4, 5, 6],
         [2, 2, 2],
         [2, 2, 2],
         [2, 2, 2]]])

如您所见,批处理中的每个内部张量都以相同的方式重复。