torch.masked#
创建时间:2022 年 8 月 15 日 | 最后更新时间:2025 年 6 月 17 日
引言#
动机#
警告
PyTorch 的掩码张量 API 处于原型阶段,将来可能会或可能不会改变。
MaskedTensor 是 torch.Tensor
的扩展,它为用户提供了以下能力:
使用任何掩码语义(例如可变长度张量、nan* 运算符等)
区分 0 和 NaN 梯度
各种稀疏应用(请参阅下面的教程)
“指定”和“未指定”在 PyTorch 中有着悠久的历史,但没有正式的语义,更没有一致性;事实上,MaskedTensor 是在香草 torch.Tensor
类无法正确解决的一系列问题之上诞生的。因此,MaskedTensor 的主要目标是成为 PyTorch 中所述“指定”和“未指定”值的真相来源,在那里它们是头等公民而不是事后诸葛亮。反过来,这应该进一步释放稀疏性的潜力,实现更安全、更一致的运算符,并为用户和开发人员提供更流畅、更直观的体验。
什么是 MaskedTensor?#
MaskedTensor 是一种张量子类,由 1) 输入(数据)和 2) 掩码组成。掩码告诉我们输入中的哪些条目应该被包含或忽略。
举个例子,假设我们想掩盖所有等于 0 的值(由灰色表示)并取最大值

上方是普通张量示例,下方是 MaskedTensor,其中所有 0 都被掩盖。这显然会根据我们是否有掩码产生不同的结果,但这种灵活的结构允许用户在计算过程中系统地忽略他们想要的任何元素。
我们已经编写了许多现有教程来帮助新用户入门,例如
支持的运算符#
一元运算符#
一元运算符是只包含单个输入的运算符。将它们应用于 MaskedTensor 相对简单:如果给定索引处的数据被掩盖,我们应用运算符,否则我们将继续掩盖数据。
可用的一元运算符包括
计算 |
|
|
|
计算 |
|
|
|
返回一个新张量,其中包含 |
|
|
|
计算给定 |
|
返回一个新张量,其中包含 |
|
|
|
返回一个新张量,其中包含 |
|
|
|
返回一个新张量,其中包含 |
|
|
|
返回一个新张量,其中包含 |
|
|
|
计算给定输入张量的按位 NOT。 |
|
返回一个新张量,其中包含 |
|
|
|
计算给定 |
|
返回一个新张量,其中包含 |
|
返回一个新张量,其中包含 |
|
返回一个新张量,其中 |
|
|
|
|
|
返回一个新张量,其元素是输入张量 |
|
|
|
|
|
返回一个新张量,其中包含 |
|
计算 |
|
计算 |
|
返回一个新张量,其中包含 |
|
返回一个新张量,其中包含 |
|
返回一个新张量,其中包含(1 + |
|
返回一个新张量,其中包含 |
|
|
|
返回一个新张量,其中包含布尔元素,表示 |
|
将 |
|
返回一个新张量,其中包含 |
|
|
|
返回 |
|
计算 |
|
返回一个新张量,其中 |
|
返回一个新张量,其中包含 |
|
将 |
|
返回一个新张量,其中包含 |
|
返回一个新张量,其中包含 |
|
此函数是 torch.sign() 对复数张量的扩展。 |
|
测试 |
|
返回一个新张量,其中包含 |
|
|
|
返回一个新张量,其中包含 |
|
返回一个新张量,其中包含 |
|
返回一个新张量,其中包含 |
|
返回一个新张量,其中包含 |
|
返回一个新张量,其元素是 |
|
返回一个新张量,其中包含 |
所有上述一元运算符都支持就地操作,**除了**
二元运算符#
如您在教程中看到的,MaskedTensor
也实现了二元操作,但需要注意的是,两个 MaskedTensor 中的掩码必须匹配,否则会引发错误。正如错误中提到的,如果您需要对特定运算符的支持,或者对它们的行为方式有建议的语义,请在 GitHub 上提出问题。目前,我们决定采用最保守的实现,以确保用户确切地了解正在发生的事情,并有意识地决定掩码语义。
可用的二元运算符包括
将 |
|
逐元素计算 的反正切,并考虑象限。 |
|
|
|
计算 |
|
计算 |
|
计算 |
|
计算 |
|
计算 |
|
将输入 |
|
|
|
逐元素应用 C++ 的 std::fmod。 |
|
输入指数和的对数。 |
|
以2为底的输入指数和的对数。 |
|
将 |
|
|
|
逐元素返回 |
|
逐元素计算 Python 的模运算。 |
|
从 |
|
|
|
当 |
|
计算逐元素相等 |
|
逐元素计算 。 |
|
逐元素计算 。 |
|
逐元素计算 。 |
|
|
|
|
|
逐元素计算 。 |
|
|
|
逐元素计算 。 |
|
|
|
计算 |
|
计算 |
|
计算 |
|
计算 |
|
|
可用的就地二元运算符包括上述所有运算符,**除了**
视图和选择函数#
我们还包含了一些视图和选择函数;直观地说,这些运算符将同时应用于数据和掩码,然后将结果封装在 MaskedTensor
中。举个简单的例子,考虑 select()
>>> data = torch.arange(12, dtype=torch.float).reshape(3, 4)
>>> data
tensor([[ 0., 1., 2., 3.],
[ 4., 5., 6., 7.],
[ 8., 9., 10., 11.]])
>>> mask = torch.tensor([[True, False, False, True], [False, True, False, False], [True, True, True, True]])
>>> mt = masked_tensor(data, mask)
>>> data.select(0, 1)
tensor([4., 5., 6., 7.])
>>> mask.select(0, 1)
tensor([False, True, False, False])
>>> mt.select(0, 1)
MaskedTensor(
[ --, 5.0000, --, --]
)
目前支持以下操作:
返回每个零维输入张量的一维视图。 |
|
根据广播语义广播给定的张量。 |
|
将 |
|
在给定维度中连接 |
|
尝试将张量分割成指定数量的块。 |
|
通过水平堆叠 |
|
根据 |
|
通过将 |
|
根据 |
|
按水平(列方向)顺序堆叠张量。 |
|
计算 |
|
创建由 attr:tensors 中一维输入指定的坐标网格。 |
|
返回一个新的张量,它是 |
|
从批量输入张量中提取滑动局部块。 |
|
返回一个连续的展平张量。 |
|
沿所选维度在给定索引处切分 |
|
将张量分割成块。 |
|
沿新维度连接一系列张量。 |
|
期望 |
|
返回一个张量,它是 |
|
根据 |
|
按垂直(行方向)顺序堆叠张量。 |
|
返回 |
|
将此张量扩展为与 |
|
返回一个张量,其数据和元素数量与 |
|
将此张量返回为与 |
|
返回原始张量的视图,其中包含 |
|
返回一个新张量,其数据与 |