python - 给定一个表示段边界的值，为张量生成段标签

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

上下文:我目前正在 PyTorch 中使用 Fairseq 的 Transformer 实现来执行 seq2seq NLP 任务。我正在寻找一种方法，在编码器的前向传递过程中将类似 BERT 的分段嵌入合并到 Transformer 中，而不是修改用于翻译任务的现有数据集，例如 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的结果。