torch.masked

创建于: Aug 15, 2022 | 最后更新于: Jun 17, 2025

简介

动机

警告

PyTorch 中带掩码张量的 API 处于原型阶段，未来可能会更改。

MaskedTensor 作为 torch.Tensor 的扩展，为用户提供了使用任何带掩码语义（例如可变长度张量、nan* 运算符等）的能力，

• 使用任何带掩码语义（例如可变长度张量、nan* 运算符等）

• 区分 0 和 NaN 梯度

• 各种稀疏应用（参见下面的教程）

“指定”和“未指定”在 PyTorch 中有着悠久的历史，但没有正式的语义，更不用说一致性了；事实上，MaskedTensor 的诞生源于 vanilla torch.Tensor 类无法妥善解决的一系列问题。因此，MaskedTensor 的首要目标是成为 PyTorch 中所述“指定”和“未指定”值的权威来源，使其成为一等公民而不是事后考虑。反过来，这应该进一步释放稀疏性的潜力，实现更安全、更一致的运算符，并为用户和开发人员提供更顺畅、更直观的体验。

什么是 MaskedTensor？

MaskedTensor 是一个张量子类，由 1) 输入（数据）和 2) 掩码组成。掩码告诉我们应包含或忽略输入中的哪些条目。

举例来说，假设我们想屏蔽所有等于 0 的值（用灰色表示）并取最大值

顶部是 vanilla 张量示例，底部是 MaskedTensor，其中所有 0 都被屏蔽。这清楚地表明结果取决于我们是否具有掩码，但这种灵活的结构允许用户在计算过程中系统地忽略他们想要的任何元素。

我们已经编写了一些现有的教程来帮助用户入门，例如

• 概述 – 新用户的起点，讨论如何使用 MaskedTensors 及其用途

• 稀疏性 – MaskedTensor 支持稀疏 COO 和 CSR 数据及掩码张量

• Adagrad 稀疏语义 – MaskedTensor 如何简化稀疏语义和实现的实用示例

• 高级语义 – 讨论为什么做出某些决定（例如，要求二进制/归约操作的掩码匹配）、与 NumPy 的 MaskedArray 的区别以及归约语义

支持的运算符

一元运算符

一元运算符是只包含一个输入的运算符。将它们应用于 MaskedTensors 相对简单：如果数据在给定索引处被屏蔽，我们会应用该运算符，否则我们将继续屏蔽数据。

可用的一元运算符是

abs	计算 input 中每个元素的绝对值。
absolute	是 torch.abs() 的别名
acos	返回一个新张量，其中包含 input 中每个元素的反余弦值（以弧度为单位）。
arccos	是 torch.acos() 的别名。
acosh	返回一个新张量，其中包含input元素的双曲反余弦。
arccosh	是 torch.acosh() 的别名。
angle	计算给定input张量的每个元素的角度（以弧度表示）。
asin	返回一个新张量，其中包含 input 张量中元素的反正弦值（以弧度为单位）。
arcsin	是 torch.asin() 的别名。
asinh	返回一个新张量，其中包含input元素的双曲反正弦。
arcsinh	是 torch.asinh() 的别名。
atan	返回一个新张量，其中包含 input 张量中元素的反正切值（以弧度为单位）。
arctan	是 torch.atan() 的别名。
atanh	返回一个新张量，其中包含input元素的双曲反正切。
arctanh	是 torch.atanh() 的别名。
bitwise_not	计算给定输入张量的按位非。
ceil	返回一个新张量，其中包含input元素的向上取整值，即大于或等于每个元素的最小整数。
clamp	将 input 中的所有元素限制在范围 [ min, max ] 内。
clip	是 torch.clamp() 的别名。
conj_physical	计算给定 input 张量的逐元素共轭。
cos	返回一个新张量，其中包含 input 中元素（以弧度为单位）的余弦值。
cosh	返回一个新张量，其中包含input元素的双曲余弦。
deg2rad	返回一个新张量，其中input的每个元素都从角度（度）转换为弧度。
digamma	是 torch.special.digamma() 的别名。
erf	是 torch.special.erf() 的别名。
erfc	是 torch.special.erfc() 的别名。
erfinv	是 torch.special.erfinv() 的别名。
exp	返回一个新张量，其元素是输入张量input的指数。
exp2	是 torch.special.exp2() 的别名。
expm1	是 torch.special.expm1() 的别名。
fix	是 torch.trunc() 的别名
floor	返回一个新张量，其中包含 input 元素的向下取整值，即小于或等于每个元素的最大整数。
frac	计算input中每个元素的小数部分。
lgamma	计算 input 上伽马函数绝对值的自然对数。
log	返回一个新张量，其中包含 input 元素对应的自然对数。
log10	返回一个新张量，其中包含input元素的以10为底的对数。
log1p	返回一个新张量，其中包含(1 + input)的自然对数。
log2	返回一个新张量，其中包含input元素的以2为底的对数。
logit	是 torch.special.logit() 的别名。
i0	是 torch.special.i0() 的别名。
isnan	返回一个新张量，其中包含布尔元素，表示 input 的每个元素是否为 NaN。
nan_to_num	分别用 nan、posinf 和 neginf 指定的值替换 input 中的 NaN、正无穷大和负无穷大。
neg	返回一个新张量，其中包含input元素的负值。
negative	是 torch.neg() 的别名
positive	返回 input。
pow	计算 input 中每个元素以 exponent 为指数的幂，并返回结果张量。
rad2deg	返回一个新张量，其中input的每个元素都从角度（弧度）转换为度。
reciprocal	返回一个新张量，其中包含input元素的倒数。
round	将input的元素四舍五入到最近的整数。
rsqrt	返回一个新张量，其中包含input每个元素的平方根的倒数。
sigmoid	是 torch.special.expit() 的别名。
sign	返回一个新张量，其中包含input元素的符号。
sgn	此函数是 torch.sign() 对复数张量的扩展。
signbit	测试input的每个元素的符号位是否已设置。
sin	返回一个新张量，其中包含 input 张量中元素（以弧度为单位）的正弦值。
sinc	是 torch.special.sinc() 的别名。
sinh	返回一个新张量，其中包含input元素的双曲正弦。
sqrt	返回一个新张量，其中包含 input 元素的平方根。
square	返回一个新张量，其中包含input元素的平方。
tan	返回一个新张量，其中包含 input 张量中元素（以弧度为单位）的正切值。
tanh	返回一个新张量，其元素是 input 的双曲正切值。
trunc	返回一个新张量，其中包含input元素的截断整数值。

可用的就地一元运算符是以上所有运算符，但不包括

angle	计算给定input张量的每个元素的角度（以弧度表示）。
positive	返回 input。
signbit	测试input的每个元素的符号位是否已设置。
isnan	返回一个新张量，其中包含布尔元素，表示 input 的每个元素是否为 NaN。

二元运算符

正如您在教程中可能看到的那样，MaskedTensor 也实现了二元运算，但前提是两个 MaskedTensor 中的掩码必须匹配，否则将引发错误。如错误所示，如果您需要对特定运算符的支持，或者对它们的行为方式有建议的语义，请在 GitHub 上提出问题。目前，我们决定采用最保守的实现，以确保用户确切地知道正在发生什么，并在带掩码语义的决策上是有意为之的。

可用的二元运算符是

add	将 other（缩放 alpha）加到 input。
atan2	在考虑象限的情况下，对 $\text{input}_{i} / \text{other}_{i}$ 进行的逐元素反正切，并考虑象限。
arctan2	是 torch.atan2() 的别名。
bitwise_and	计算 input 和 other 的按位与。
bitwise_or	计算 input 和 other 的按位或。
bitwise_xor	计算 input 和 other 的按位异或。
bitwise_left_shift	计算 input 按 other 位进行左算术移位。
bitwise_right_shift	计算 input 按 other 位进行右算术移位。
div	将输入input的每个元素除以other的相应元素。
divide	是 torch.div() 的别名。
floor_divide
fmod	逐元素应用 C++ 的 std::fmod。
logaddexp	输入指数和的对数。
logaddexp2	以2为底的输入指数和的对数。
mul	将input乘以other。
multiply	是 torch.mul() 的别名。
nextafter	返回 input 之后，趋向于 other 的下一个浮点值，逐元素进行。
remainder	逐元素计算 Python 的模数运算。
sub	从 input 中减去 other（缩放 alpha）。
subtract	是 torch.sub() 的别名。
true_divide	是 torch.div() 的别名，其中 rounding_mode=None。
eq	计算逐元素相等
ne	逐元素计算 $\text{input} \neq \text{other}$。
le	逐元素计算 $\text{input} \leq \text{other}$。
ge	逐元素计算 $\text{input} \geq \text{other}$。
greater	是 torch.gt() 的别名。
greater_equal	是 torch.ge() 的别名。
gt	逐元素计算 $\text{input} > \text{other}$。
less_equal	是 torch.le() 的别名。
lt	逐元素计算 $\text{input} < \text{other}$。
less	是 torch.lt() 的别名。
maximum	计算 input 和 other 的逐元素最大值。
minimum	计算 input 和 other 的逐元素最小值。
fmax	计算 input 和 other 的逐元素最大值。
fmin	计算 input 和 other 的逐元素最小值。
not_equal	是 torch.ne() 的别名。

可用的就地二元运算符是以上所有运算符，但不包括

logaddexp	输入指数和的对数。
logaddexp2	以2为底的输入指数和的对数。
equal	如果两个张量具有相同的大小和元素，则为True，否则为False。
fmin	计算 input 和 other 的逐元素最小值。
minimum	计算 input 和 other 的逐元素最小值。
fmax	计算 input 和 other 的逐元素最大值。

归约

可用的归约操作如下（支持自动微分）。有关更多信息，概述教程介绍了一些归约示例，而高级语义教程对我们如何确定某些归约语义进行了更深入的讨论。

sum	返回 input 张量中所有元素的和。
mean
amin	在给定维度 dim 下，返回 input 张量每个切片的最小值。
amax	在给定维度 dim 下，返回 input 张量每个切片的最大值。
argmin	返回扁平张量或沿某一维度的最小值的索引
argmax	返回 input 张量中所有元素最大值的索引。
prod	返回input张量中所有元素的乘积。
all	测试 input 中的所有元素是否都评估为 True。
norm	返回给定张量的矩阵范数或向量范数。
var	计算在由 dim 指定的维度上的方差。
std	计算在由 dim 指定的维度上的标准差。

视图和选择函数

我们还包含了一些视图和选择函数；直观地说，这些运算符将应用于数据和掩码，然后将结果包装在 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,       --,       --]
    )

当前支持以下操作

atleast_1d	返回每个输入张量的一维视图，其中零维度。
broadcast_tensors	根据 广播语义 广播给定的张量。
broadcast_to	将 input 广播到形状 shape。
cat	沿给定维度连接 tensors 中的给定张量序列。
chunk	尝试将张量分割成指定数量的块。
column_stack	通过水平堆叠tensors中的张量来创建新张量。
dsplit	根据 indices_or_sections，沿深度方向将 input（一个具有三层或更多维度的张量）分割成多个张量。
flatten	通过将 input 重塑为一维张量来展平。
hsplit	根据 indices_or_sections，沿水平方向将 input（一个具有一维或更多维度的张量）分割成多个张量。
hstack	按水平（列方向）顺序堆叠张量。
kron	计算 input 和 other 的克罗内克积，表示为 $\otimes$。
meshgrid	创建由 attr:tensors 中一维输入指定的坐标网格。
narrow	返回 input 张量的缩小版本。
nn.functional.unfold	从批量输入张量中提取滑动局部块。
ravel	返回一个连续的展平张量。
select	在选定维度上的给定索引处切片 input 张量。
split	将张量分割成块。
stack	沿新维度连接一系列张量。
t	期望input为小于等于2维的张量，并转置维度0和1。
转置	返回 input 张量的转置版本。
vsplit	根据 indices_or_sections，沿垂直方向将 input（一个具有二维或更多维度的张量）分割成多个张量。
vstack	按垂直（行方向）顺序堆叠张量。
Tensor.expand	返回 self 张量的新视图，其中单例维度已扩展到更大的大小。
Tensor.expand_as	将此张量扩展到与 other 相同的尺寸。
Tensor.reshape	返回一个与 self 具有相同数据和元素数量的张量，但具有指定的形状。
Tensor.reshape_as	将此张量以与 other 相同的形状返回。
Tensor.unfold	返回原始张量的视图，其中包含 self 张量在 dimension 维度上所有大小为 size 的切片。
Tensor.view	返回一个新张量，其数据与 self 张量相同，但 shape 不同。