python - Pytorch:将值从一个掩码分配给另一个掩码，由自身掩码

Question

我有一个掩码 active 用于跟踪在循环过程中仍未终止的批处理。它的维度是[batch_full,]，它的真实条目显示哪些元素在当前步骤中仍需要使用。循环过程生成另一个掩码 terminated，它具有与 active 掩码中的真值一样多的元素。现在，我想从 ~terminated 中获取值并将它们放回到 active 中，但在正确的索引处。基本上我想做的是:

import torch

active = torch.ones([4,], dtype=torch.bool)
active[:2] = torch.tensor(False)

terminated = torch.tensor([True, False])

active[active] = ~terminated

print(active)  # expected [F, F, F, T]

但是，我得到错误:

RuntimeError: unsupported operation: some elements of the input tensor and the written-to tensor refer to a single memory location. Please clone() the tensor before performing the operation.

如何有效地进行上述操作？

Answer 1

有几个解决方案，我也会给出它们的速度，通过 timeit 测量，重复 10k，在 2021 macbook pro 上。

最简单的解决方案，耗时 0.260 秒:

active[active.clone()] = ~terminated

我们可以对 abt 使用 masked_scatter_ 就地操作。 2 倍加速(0.136 秒):

active.masked_scatter_(
        active,
        ~terminated,
    )

错位操作，耗时 0.161 秒，将是:

active = torch.masked_scatter(
        active,
        active,
        ~terminated,
    )