PyTorch 使用布尔掩码（保留维度）提取张量元素

Question

比如说，我有一个 PyTorch 2x2 张量，我还生成了一个相同维度 (2x2) 的布尔张量。我想用这个做面膜

例如，如果我有一个张量：

tensor([[1, 3],
        [4, 7]])

如果我的面具是：

tensor([[ True, False],
        [False,  True]])

我想使用该掩码获取一个张量，其中保留与原始张量 True 相对应的元素，而在输出张量中将与 False 相对应的元素设置为零。

预期输出：

tensor([[1, 0],
        [0, 7]])

感谢任何帮助。谢谢！

Answer 1

假设你有：

t = torch.Tensor([[1,2], [3,4]])
mask = torch.Tensor([[True,False], [False,True]])

您可以通过以下方式使用面具：

masked_t = t * mask

输出将是：

tensor([[1., 0.],
        [0., 4.]])

PyTorch extracting tensor elements with boolean mask (retaining dimensions)