给定一个指示段边界的值，为 Tensor 生成段标签

Question

有谁知道为张量生成 'segment label' 的方法，给定一个表示张量内段边界的唯一值？

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

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

生成的段标签 Tensor 应具有相同的形状，其值代表两个段：

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 的 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。

给定一个指示段边界的值，为 Tensor 生成段标签

generating segment labels for a Tensor given a value indicating segment boundaries

python

pytorch

tensor

bert-language-model

fairseq