python - pytorch 在张量的某个维度上调用 unique_consecutive() 的最佳方法？

Question

在张量的特定维度上调用 unique_consecutive() 并用指定​​值填充剩余部分的最佳方法是什么？为简单起见，我们可以使用 2D 张量作为示例:

input = tensor([[3, 3, 5, 5, 5],
                [3, 3, 2, 2, 3]])

如果指定padding值为-1，我希望得到的是:

output = tensor([[3, 5, -1, -1, -1],
                 [3, 2, 3, -1, -1]])

此外，为了节省内存，如果可能的话，我更愿意就地这样做。谢谢!

Answer 1

从 PyTorch 论坛找到了一个不错的解决方案 https://discuss.pytorch.org/t/best-way-to-run-unique-consecutive-on-certain-dimension/112662 ，此处引用以供快速引用。

    x = tensor([[3, 3, 5, 5, 5],
                [3, 3, 2, 2, 3]])
    unique_x, indices = torch.unique_consecutive(x, return_inverse=True)
    indices -= indices.min(dim=1, keepdims=True)[0]
    result = -torch.ones_like(x)
    result = result.scatter_(1, indices, x)