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

Question

有人知道为张量生成"分段标签"的方法吗？给定一个表示张量内分段边界的唯一值？

例如，给定1D输入张量，其中值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]
])

上下文：我目前正在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的结果。