当前位置：首页 > news >正文

pytorch中不同的mask方法：masked_fill, masked_select, masked_scatter

news 来源：原创 2025/6/22 15:34:23

在 PyTorch 中，masked_fill、masked_select 和 masked_scatter 是三种常用的掩码（mask）操作方法，它们通过布尔类型的掩码张量（mask）对原始张量进行条件筛选或修改。以下是它们的详细解释和对比：

1. masked_fill

作用：将原始张量中 mask 为 True 的位置用指定值填充，其余位置保持不变。

参数：
• mask（BoolTensor）：与原始张量形状相同的布尔掩码。
• value（标量）：用于填充的值。

特点：
• 原地操作：会直接修改原始张量（除非使用 masked_fill_ 的 in-place 版本）。
• 保持形状：输出张量形状与输入张量一致。

示例：

import torch

x = torch.tensor([[1, 2], [3, 4]])
mask = torch.tensor([[False, True], [True, False]])

y = x.masked_fill(mask, -1)
print(y)
# 输出：
# tensor([[ 1, -1],
#         [-1,  4]])

典型应用：
• 在 Transformer 的注意力机制中，用 -inf 填充 padding 或未来的位置，使 softmax 后概率为 0。
• 数据清洗时屏蔽无效值（如 NaN）。

2. masked_select

作用：从原始张量中提取 mask 为 True 的元素，返回一维张量。

参数：
• mask（BoolTensor）：与原始张量形状相同的布尔掩码。

特点：
• 返回一维张量：输出会丢失原始张量的维度信息。
• 非原地操作：生成新的张量。

示例：

x = torch.tensor([[1, 2], [3, 4]])
mask = torch.tensor([[False, True], [True, False]])

y = x.masked_select(mask)
print(y)  # tensor([2, 3])

典型应用：
• 提取满足条件的元素（如分类任务中筛选正样本）。
• 统计掩码区域的值（如计算非零元素均值）。

3. masked_scatter

作用：将另一个张量（source）中的值按顺序填充到原始张量中 mask 为 True 的位置。

参数：
• mask（BoolTensor）：与原始张量形状相同的布尔掩码。
• source（Tensor）：提供填充值的源张量。

特点：
• 按顺序填充：source 中的值按行优先顺序填充到 mask 为 True 的位置。
• source 的长度必须 ≥ mask 中 True 的数量。

示例：

x = torch.tensor([[1, 2], [3, 4]])
mask = torch.tensor([[False, True], [True, False]])
source = torch.tensor([10, 20])

y = x.masked_scatter(mask, source)
print(y)
# 输出：
# tensor([[ 1, 10],
#         [20,  4]])

典型应用：
• 动态替换张量中的部分值（如用随机噪声替换特定区域）。
• 批量更新参数时选择性地替换某些位置。

对比总结

方法	输入张量形状	输出形状	是否修改原张量	核心功能
`masked_fill`	保留原形状	与原张量相同	是（可选）	用标量填充掩码区域
`masked_select`	保留原形状	一维张量	否	提取掩码区域的元素
`masked_scatter`	保留原形状	与原张量相同	是（可选）	用另一张量填充掩码区域