PyTorch 使用布尔掩码提取张量元素（保留尺寸）。

Question

比如说，我有一个 PyTorch 2x2 张量，我还生成了一个相同维度 (2x2) 的布尔张量。我想把它作为一个掩码。

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

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

如果我的掩模是：

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

我想用这个掩码来得到一个张量，其中对应的元素是 True 保留了我的原始张量，而在输出张量中，对应于False的元素被设置为零。

预期的输出。

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

希望能得到帮助 谢谢！我有一个2x2的PyTorch张量，我也有一个2x2的张量。

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.]])