为给定的值表示段边界的张量生成段标签

Question

有没有人知道为张量生成“段标签”的方法，给出一个表示张量内段边界的唯一值？

例如，给定一个一维输入张量，其中值1表示分段边界，

x = torch.Tensor([5, 4, 1, 3, 6, 2])

由此产生的段标签张量应该具有相同的形状，其值表示两个段：

segment_label = torch.Tensor([1, 1, 1, 2, 2, 2])

同样，对于一批输入，例如批大小= 3，

x = torch.Tensor([
    [5, 4, 1, 3, 6, 2],
    [9, 4, 5, 1, 8, 10],
    [10, 1, 5, 4, 8, 9]
    ])

生成的段标签张量(使用1作为段分隔符)应该如下所示：

segment_label = torch.Tensor([
    [1, 1, 1, 2, 2, 2],
    [1, 1, 1, 1, 2, 2],
    [1, 1, 2, 2, 2, 2]
    ])

Context：我目前正在PyTorch中使用Fairseq的变压器实现来执行seq2seq NLP任务。我正在寻找一种方法，以便在编码器的前向传递过程中将类似于伯特的段嵌入合并到变压器中，而不是修改用于翻译任务(如language_pair_dataset )的现有数据集。

提前感谢！

Answer 1

您可以使用torch.cumsum来完成以下操作：

mask = (x == 1).to(x)  # mask with only the boundaries
segment_label = mask.cumsum(dim=-1) - mask + 1

结果与预期的segment_label。