python - 用掩码为真的另一个张量填充张量

Question

我需要以一定的概率将张量 new 的元素插入到张量 old 中，为简单起见，我们假设它是 0.8。基本上这就是 masked_fill 会做的，但它只适用于一维张量。其实我在做

    prob = torch.rand(trgs.shape, dtype=torch.float32).to(trgs.device)
    mask = prob < 0.8

    dim1, dim2, dim3, dim4 = new.shape
    for a in range(dim1):
        for b in range(dim2):
            for c in range(dim3):
                for d in range(dim4):
                    old[a][b][c][d] = old[a][b][c][d] if mask[a][b][c][d] else new[a][b][c][d]

这太糟糕了。我想要类似的东西

    prob = torch.rand(trgs.shape, dtype=torch.float32).to(trgs.device)
    mask = prob < 0.8

    old = trgs.multidimensional_masked_fill(mask, new)

Answer 1

我不确定你的一些对象是什么，但这应该能让你在短时间内完成你需要做的事情:

old 是您现有的数据。

mask 是您以概率 p 生成的掩码

new 是包含您要插入的元素的新张量。

# torch.where
result = old.where(mask, new)