评价此页
★ ★ ★ ★ ★
unstable/maskedtensor_overview

注意

跳转至页面底部 下载完整示例代码。

MaskedTensor 概述#

本教程旨在作为使用 MaskedTensor 的起点，并讨论其掩码语义。

MaskedTensor 是 torch.Tensor 的扩展，它为用户提供了以下能力：

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

  • 区分 0 梯度和 NaN 梯度

  • 各种稀疏应用（见下方教程）

有关 MaskedTensor 的更详细介绍，请参阅 torch.masked 文档。

使用 MaskedTensor#

在本节中，我们将讨论如何使用 MaskedTensor，包括如何构建、访问数据和掩码，以及索引和切片操作。

准备工作#

我们首先进行教程所需的环境设置

import torch
from torch.masked import masked_tensor, as_masked_tensor
import warnings

# Disable prototype warnings and such
warnings.filterwarnings(action='ignore', category=UserWarning)

构造#

构建 MaskedTensor 有几种不同的方法

  • 第一种方法是直接调用 MaskedTensor 类

  • 第二种（也是我们推荐的）方法是使用 masked.masked_tensor() 和 masked.as_masked_tensor() 工厂函数，它们类似于 torch.tensor() 和 torch.as_tensor()

在本教程中，我们假设已执行导入语句：from torch.masked import masked_tensor。

访问数据和掩码#

MaskedTensor 中的底层字段可以通过以下方式访问：

  • MaskedTensor.get_data() 函数

  • MaskedTensor.get_mask() 函数。请记住，True 表示“已指定”或“有效”，而 False 表示“未指定”或“无效”。

通常，返回的底层数据在未指定的条目中可能无效，因此我们建议当用户需要不包含任何被掩码条目的张量时，使用 MaskedTensor.to_tensor()（如上所示）来返回一个带有填充值的张量。

索引和切片#

MaskedTensor 是一个 Tensor 子类，这意味着它继承了与 torch.Tensor 相同的索引和切片语义。以下是一些常见索引和切片模式的示例

data = torch.arange(24).reshape(2, 3, 4)
mask = data % 2 == 0

print("data:\n", data)
print("mask:\n", mask)

data:
 tensor([[[ 0,  1,  2,  3],
         [ 4,  5,  6,  7],
         [ 8,  9, 10, 11]],

        [[12, 13, 14, 15],
         [16, 17, 18, 19],
         [20, 21, 22, 23]]])
mask:
 tensor([[[ True, False,  True, False],
         [ True, False,  True, False],
         [ True, False,  True, False]],

        [[ True, False,  True, False],
         [ True, False,  True, False],
         [ True, False,  True, False]]])

# float is used for cleaner visualization when being printed
mt = masked_tensor(data.float(), mask)

print("mt[0]:\n", mt[0])
print("mt[:, :, 2:4]:\n", mt[:, :, 2:4])

mt[0]:
 MaskedTensor(
  [
    [  0.0000,       --,   2.0000,       --],
    [  4.0000,       --,   6.0000,       --],
    [  8.0000,       --,  10.0000,       --]
  ]
)
mt[:, :, 2:4]:
 MaskedTensor(
  [
    [
      [  2.0000,       --],
      [  6.0000,       --],
      [ 10.0000,       --]
    ],
    [
      [ 14.0000,       --],
      [ 18.0000,       --],
      [ 22.0000,       --]
    ]
  ]
)

为什么 MaskedTensor 很有用？#

由于 MaskedTensor 将已指定和未指定的值视为一等公民，而不是事后处理（如使用填充值、NaN 等），因此它能够解决常规张量无法解决的几个缺陷；实际上，MaskedTensor 的诞生在很大程度上正是由于这些反复出现的问题。

下面，我们将讨论目前 PyTorch 中一些尚未解决的常见问题，并说明 MaskedTensor 如何解决这些问题。

区分 0 和 NaN 梯度#

torch.Tensor 遇到的一个问题是无法区分未定义的梯度 (NaN) 和实际为 0 的梯度。由于 PyTorch 没有标记值是“已指定/有效”还是“未指定/无效”的方法，它被迫依赖 NaN 或 0（取决于具体用例），这导致了不可靠的语义，因为许多操作无法正确处理 NaN 值。更令人困惑的是，有时根据操作顺序的不同，梯度可能会发生变化（例如，取决于 NaN 值在操作链中出现的早晚）。

MaskedTensor 是解决这个问题的完美方案！

torch.where#

在 Issue 10729 中，我们注意到在使用 torch.where() 时，操作顺序可能会产生影响，因为我们难以区分 0 是真实的 0 还是来自未定义梯度的 0。因此，为了保持一致性，我们选择掩盖掉结果。

当前结果

x = torch.tensor([-10., -5, 0, 5, 10, 50, 60, 70, 80, 90, 100], requires_grad=True, dtype=torch.float)
y = torch.where(x < 0, torch.exp(x), torch.ones_like(x))
y.sum().backward()
x.grad

tensor([4.5400e-05, 6.7379e-03, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
        0.0000e+00, 0.0000e+00, 0.0000e+00,        nan,        nan])

MaskedTensor 结果

x = torch.tensor([-10., -5, 0, 5, 10, 50, 60, 70, 80, 90, 100])
mask = x < 0
mx = masked_tensor(x, mask, requires_grad=True)
my = masked_tensor(torch.ones_like(x), ~mask, requires_grad=True)
y = torch.where(mask, torch.exp(mx), my)
y.sum().backward()
mx.grad

MaskedTensor(
  [  0.0000,   0.0067,       --,       --,       --,       --,       --,       --,       --,       --,       --]
)

这里的梯度仅提供给选定的子集。实际上，这改变了 where 的梯度，使其掩盖掉元素而不是将它们设为零。

另一个 torch.where#

Issue 52248 是另一个示例。

当前结果

a = torch.randn((), requires_grad=True)
b = torch.tensor(False)
c = torch.ones(())
print("torch.where(b, a/0, c):\n", torch.where(b, a/0, c))
print("torch.autograd.grad(torch.where(b, a/0, c), a):\n", torch.autograd.grad(torch.where(b, a/0, c), a))

torch.where(b, a/0, c):
 tensor(1., grad_fn=<WhereBackward0>)
torch.autograd.grad(torch.where(b, a/0, c), a):
 (tensor(nan),)

MaskedTensor 结果

a = masked_tensor(torch.randn(()), torch.tensor(True), requires_grad=True)
b = torch.tensor(False)
c = torch.ones(())
print("torch.where(b, a/0, c):\n", torch.where(b, a/0, c))
print("torch.autograd.grad(torch.where(b, a/0, c), a):\n", torch.autograd.grad(torch.where(b, a/0, c), a))

torch.where(b, a/0, c):
 MaskedTensor(  1.0000, True)
torch.autograd.grad(torch.where(b, a/0, c), a):
 (MaskedTensor(--, False),)

这个问题很类似（甚至链接到了下文的下一个问题），它表达了对意外行为的挫败感，因为无法区分“无梯度”和“零梯度”，这反过来使得处理其他操作变得难以推断。

使用掩码时，x/0 会产生 NaN 梯度#

在 Issue 4132 中，用户建议 x.grad 应该是 [0, 1] 而不是 [nan, 1]，而 MaskedTensor 通过完全掩盖掉梯度使得这一点非常明确。

当前结果

x = torch.tensor([1., 1.], requires_grad=True)
div = torch.tensor([0., 1.])
y = x/div # => y is [inf, 1]
mask = (div != 0)  # => mask is [0, 1]
y[mask].backward()
x.grad

tensor([nan, 1.])

MaskedTensor 结果

x = torch.tensor([1., 1.], requires_grad=True)
div = torch.tensor([0., 1.])
y = x/div # => y is [inf, 1]
mask = (div != 0) # => mask is [0, 1]
loss = as_masked_tensor(y, mask)
loss.sum().backward()
x.grad

MaskedTensor(
  [      --,   1.0000]
)

torch.nansum() 和 torch.nanmean()#

在 Issue 67180 中，梯度计算不正确（这是一个长期存在的问题），而 MaskedTensor 则能正确处理。

当前结果

a = torch.tensor([1., 2., float('nan')])
b = torch.tensor(1.0, requires_grad=True)
c = a * b
c1 = torch.nansum(c)
bgrad1, = torch.autograd.grad(c1, b, retain_graph=True)
bgrad1

tensor(nan)

MaskedTensor 结果

a = torch.tensor([1., 2., float('nan')])
b = torch.tensor(1.0, requires_grad=True)
mt = masked_tensor(a, ~torch.isnan(a))
c = mt * b
c1 = torch.sum(c)
bgrad1, = torch.autograd.grad(c1, b, retain_graph=True)
bgrad1

MaskedTensor(  3.0000, True)

安全 Softmax#

安全 Softmax 是另一个经常出现的问题的极佳示例。简而言之，如果有一个完整的批次被“掩盖”或完全由填充组成（在 softmax 的情况下，转化为被设为 -inf），那么这将导致 NaN，从而可能导致训练发散。

幸运的是，MaskedTensor 已经解决了这个问题。考虑这种设置

data = torch.randn(3, 3)
mask = torch.tensor([[True, False, False], [True, False, True], [False, False, False]])
x = data.masked_fill(~mask, float('-inf'))
mt = masked_tensor(data, mask)
print("x:\n", x)
print("mt:\n", mt)

x:
 tensor([[-0.5788,    -inf,    -inf],
        [ 2.0319,    -inf, -0.4221],
        [   -inf,    -inf,    -inf]])
mt:
 MaskedTensor(
  [
    [ -0.5788,       --,       --],
    [  2.0319,       --,  -0.4221],
    [      --,       --,       --]
  ]
)

例如，我们要计算沿 dim=0 的 softmax。请注意，第二列是“不安全”的（即完全被掩盖），因此当计算 softmax 时，结果将产生 0/0 = nan，因为 exp(-inf) = 0。然而，我们真正想要的是梯度被掩盖掉，因为它们未被指定，对训练来说是无效的。

PyTorch 结果

x.softmax(0)

tensor([[0.0685,    nan, 0.0000],
        [0.9315,    nan, 1.0000],
        [0.0000,    nan, 0.0000]])

MaskedTensor 结果

mt.softmax(0)

MaskedTensor(
  [
    [  0.0685,       --,       --],
    [  0.9315,       --,   1.0000],
    [      --,       --,       --]
  ]
)

实现缺失的 torch.nan* 运算符#

在 Issue 61474 中，有人请求添加额外的运算符以涵盖各种 torch.nan* 应用，例如 torch.nanmax、torch.nanmin 等。

总的来说，这些问题更自然地适用于掩码语义，因此我们建议使用 MaskedTensor 代替引入额外的运算符。由于 nanmean 已经实现，我们可以用它作为比较点

x = torch.arange(16).float()
y = x * x.fmod(4)
z = y.masked_fill(y == 0, float('nan'))  # we want to get the mean of y when ignoring the zeros

print("y:\n", y)
# z is just y with the zeros replaced with nan's
print("z:\n", z)

y:
 tensor([ 0.,  1.,  4.,  9.,  0.,  5., 12., 21.,  0.,  9., 20., 33.,  0., 13.,
        28., 45.])
z:
 tensor([nan,  1.,  4.,  9., nan,  5., 12., 21., nan,  9., 20., 33., nan, 13.,
        28., 45.])

print("y.mean():\n", y.mean())
print("z.nanmean():\n", z.nanmean())
# MaskedTensor successfully ignores the 0's
print("torch.mean(masked_tensor(y, y != 0)):\n", torch.mean(masked_tensor(y, y != 0)))

y.mean():
 tensor(12.5000)
z.nanmean():
 tensor(16.6667)
torch.mean(masked_tensor(y, y != 0)):
 MaskedTensor( 16.6667, True)

在上面的示例中，我们构建了一个 y 并希望在忽略零的情况下计算序列的平均值。torch.nanmean 可以用来实现这一点，但我们没有实现其余的 torch.nan* 操作。MaskedTensor 通过能够使用基础操作解决了这个问题，并且我们已经支持了 issue 中列出的其他操作。例如：

torch.argmin(masked_tensor(y, y != 0))

MaskedTensor(  1.0000, True)

事实上，在忽略 0 时，最小参数的索引是索引 1 处的 1。

MaskedTensor 还可以在数据完全被掩盖时支持约简操作，这等同于上述数据张量完全为 nan 的情况。nanmean 会返回 nan（一个模糊的返回值），而 MaskedTensor 则能更准确地指示出一个被掩盖的结果。 

x = torch.empty(16).fill_(float('nan'))
print("x:\n", x)
print("torch.nanmean(x):\n", torch.nanmean(x))
print("torch.nanmean via maskedtensor:\n", torch.mean(masked_tensor(x, ~torch.isnan(x))))

x:
 tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan])
torch.nanmean(x):
 tensor(nan)
torch.nanmean via maskedtensor:
 MaskedTensor(--, False)

这是一个与安全 softmax 类似的问题，即当我们需要一个未定义值时，0/0 = nan。

结论#

在本教程中，我们介绍了什么是 MaskedTensor，展示了如何使用它们，并通过一系列它们所解决的问题和示例激发了它们的价值。

进一步阅读#

要继续深入了解，您可以查看我们的 MaskedTensor 稀疏性教程，了解 MaskedTensor 如何实现稀疏性以及我们当前支持的不同存储格式。

脚本运行总耗时：（0 分 0.156 秒）

在此页面上
PyTorch 库