（原型）MaskedTensor 稀疏性

创建于：2022 年 10 月 28 日 | 最后更新：2023 年 12 月 12 日 | 最后验证：未验证

在学习本教程之前，请务必查看我们的 MaskedTensor 概述教程 <https://pytorch.ac.cn/tutorials/prototype/maskedtensor_overview.html>。

简介

稀疏性一直是 PyTorch 中快速增长且重要的领域；如果以下任何稀疏性术语让您感到困惑，请参考稀疏性教程以获取更多详细信息。

稀疏存储格式已被证明在多种方面都非常强大。作为入门，大多数从业者首先想到的用例是当大多数元素都等于零（高度稀疏）时，但即使在稀疏度较低的情况下，某些格式（例如 BSR）也可以利用矩阵内的子结构。

注意

目前，MaskedTensor 支持 COO 和 CSR 张量，并计划在未来支持其他格式（例如 BSR 和 CSC）。如果您对其他格式有任何要求，请在此处提交功能请求 here！

原则

在创建具有稀疏张量的 MaskedTensor 时，必须遵守以下几个原则

1. data 和 mask 必须具有相同的存储格式，无论是 torch.strided、torch.sparse_coo 还是 torch.sparse_csr

2. data 和 mask 必须具有相同的大小，由 size() 指示

稀疏 COO 张量

根据原则 #1，稀疏 COO MaskedTensor 是通过传入两个稀疏 COO 张量创建的，这两个张量可以通过其任何构造函数初始化，例如 torch.sparse_coo_tensor()。

作为 稀疏 COO 张量 的回顾，COO 格式代表“坐标格式”，其中指定的元素存储为其索引和相应值的元组。也就是说，提供以下内容

• indices：大小为 (ndim, nse) 且 dtype 为 torch.int64 的数组

• values：大小为 (nse,) 且具有任何整数或浮点 dtype 的数组

其中 ndim 是张量的维度，nse 是指定元素的数量。

对于稀疏 COO 和 CSR 张量，您都可以通过以下两种方式构造 MaskedTensor

1. masked_tensor(sparse_tensor_data, sparse_tensor_mask)

2. dense_masked_tensor.to_sparse_coo() 或 dense_masked_tensor.to_sparse_csr()

第二种方法更容易说明，因此我们在下面展示了它，但有关第一种方法以及方法背后的细微差别，请阅读 稀疏 COO 附录。

import torch
from torch.masked import masked_tensor
import warnings

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

values = torch.tensor([[0, 0, 3], [4, 0, 5]])
mask = torch.tensor([[False, False, True], [False, False, True]])
mt = masked_tensor(values, mask)
sparse_coo_mt = mt.to_sparse_coo()

print("mt:\n", mt)
print("mt (sparse coo):\n", sparse_coo_mt)
print("mt data (sparse coo):\n", sparse_coo_mt.get_data())

mt:
 MaskedTensor(
  [
    [      --,       --, 3],
    [      --,       --, 5]
  ]
)
mt (sparse coo):
 MaskedTensor(
  [
    [      --,       --, 3],
    [      --,       --, 5]
  ]
)
mt data (sparse coo):
 tensor(indices=tensor([[0, 1],
                       [2, 2]]),
       values=tensor([3, 5]),
       size=(2, 3), nnz=2, layout=torch.sparse_coo)

稀疏 CSR 张量

同样，MaskedTensor 也支持 CSR（压缩稀疏行） 稀疏张量格式。稀疏 CSR 张量不存储像稀疏 COO 张量那样的索引元组，而是旨在通过存储压缩行索引来减少内存需求。特别是，CSR 稀疏张量由三个 1-D 张量组成

• crow_indices：压缩行索引数组，大小为 (size[0] + 1,)。此数组指示给定值条目位于哪一行。最后一个元素是指定元素的数量，而 crow_indices[i+1] - crow_indices[i] 指示第 i 行中指定元素的数量。

• col_indices：大小为 (nnz,) 的数组。指示每个值的列索引。

• values：大小为 (nnz,) 的数组。包含 CSR 张量的值。

值得注意的是，稀疏 COO 和 CSR 张量都处于 beta 状态。

举例说明

mt_sparse_csr = mt.to_sparse_csr()

print("mt (sparse csr):\n", mt_sparse_csr)
print("mt data (sparse csr):\n", mt_sparse_csr.get_data())

mt (sparse csr):
 MaskedTensor(
  [
    [      --,       --, 3],
    [      --,       --, 5]
  ]
)
mt data (sparse csr):
 tensor(crow_indices=tensor([0, 1, 2]),
       col_indices=tensor([2, 2]),
       values=tensor([3, 5]), size=(2, 3), nnz=2, layout=torch.sparse_csr)

支持的操作

一元运算

支持所有 一元运算符，例如

mt.sin()

MaskedTensor(
  [
    [      --,       --,   0.1411],
    [      --,       --,  -0.9589]
  ]
)

二元运算

也支持 二元运算符，但来自两个 MaskedTensor 的输入掩码必须匹配。有关为何做出此决定的更多信息，请参阅我们的 MaskedTensor：高级语义教程。

请在下面找到示例

i = [[0, 1, 1],
     [2, 0, 2]]
v1 = [3, 4, 5]
v2 = [20, 30, 40]
m = torch.tensor([True, False, True])

s1 = torch.sparse_coo_tensor(i, v1, (2, 3))
s2 = torch.sparse_coo_tensor(i, v2, (2, 3))
mask = torch.sparse_coo_tensor(i, m, (2, 3))

mt1 = masked_tensor(s1, mask)
mt2 = masked_tensor(s2, mask)

print("mt1:\n", mt1)
print("mt2:\n", mt2)

mt1:
 MaskedTensor(
  [
    [      --,       --, 3],
    [      --,       --, 5]
  ]
)
mt2:
 MaskedTensor(
  [
    [      --,       --, 20],
    [      --,       --, 40]
  ]
)

print("torch.div(mt2, mt1):\n", torch.div(mt2, mt1))
print("torch.mul(mt1, mt2):\n", torch.mul(mt1, mt2))

torch.div(mt2, mt1):
 MaskedTensor(
  [
    [      --,       --,   6.6667],
    [      --,       --,   8.0000]
  ]
)
torch.mul(mt1, mt2):
 MaskedTensor(
  [
    [      --,       --, 60],
    [      --,       --, 200]
  ]
)

归约运算

最后，支持 归约运算

mt

MaskedTensor(
  [
    [      --,       --, 3],
    [      --,       --, 5]
  ]
)

print("mt.sum():\n", mt.sum())
print("mt.sum(dim=1):\n", mt.sum(dim=1))
print("mt.amin():\n", mt.amin())

mt.sum():
 MaskedTensor(8, True)
mt.sum(dim=1):
 MaskedTensor(
  [3, 5]
)
mt.amin():
 MaskedTensor(3, True)

MaskedTensor 辅助方法

为了方便起见，MaskedTensor 有许多方法可以帮助在不同布局之间进行转换并识别当前布局

设置

v = [[3, 0, 0],
     [0, 4, 5]]
m = [[True, False, False],
     [False, True, True]]

mt = masked_tensor(torch.tensor(v), torch.tensor(m))
mt

MaskedTensor(
  [
    [3,       --,       --],
    [      --, 4, 5]
  ]
)

MaskedTensor.to_sparse_coo() / MaskedTensor.to_sparse_csr() / MaskedTensor.to_dense() 帮助在不同布局之间进行转换。

mt_sparse_coo = mt.to_sparse_coo()
mt_sparse_csr = mt.to_sparse_csr()
mt_dense = mt_sparse_coo.to_dense()

MaskedTensor.is_sparse() – 这将检查 MaskedTensor 的布局是否与任何受支持的稀疏布局（目前为 COO 和 CSR）匹配。

print("mt_dense.is_sparse: ", mt_dense.is_sparse)
print("mt_sparse_coo.is_sparse: ", mt_sparse_coo.is_sparse)
print("mt_sparse_csr.is_sparse: ", mt_sparse_csr.is_sparse)

mt_dense.is_sparse:  False
mt_sparse_coo.is_sparse:  True
mt_sparse_csr.is_sparse:  True

MaskedTensor.is_sparse_coo()

print("mt_dense.is_sparse_coo(): ", mt_dense.is_sparse_coo())
print("mt_sparse_coo.is_sparse_coo: ", mt_sparse_coo.is_sparse_coo())
print("mt_sparse_csr.is_sparse_coo: ", mt_sparse_csr.is_sparse_coo())

mt_dense.is_sparse_coo():  False
mt_sparse_coo.is_sparse_coo:  True
mt_sparse_csr.is_sparse_coo:  False

MaskedTensor.is_sparse_csr()

print("mt_dense.is_sparse_csr(): ", mt_dense.is_sparse_csr())
print("mt_sparse_coo.is_sparse_csr: ", mt_sparse_coo.is_sparse_csr())
print("mt_sparse_csr.is_sparse_csr: ", mt_sparse_csr.is_sparse_csr())

mt_dense.is_sparse_csr():  False
mt_sparse_coo.is_sparse_csr:  False
mt_sparse_csr.is_sparse_csr:  True

附录

稀疏 COO 构造

回想一下我们在 原始示例 中，我们创建了一个 MaskedTensor，然后使用 MaskedTensor.to_sparse_coo() 将其转换为稀疏 COO MaskedTensor。

或者，我们也可以通过传入两个稀疏 COO 张量直接构造稀疏 COO MaskedTensor

values = torch.tensor([[0, 0, 3], [4, 0, 5]]).to_sparse()
mask = torch.tensor([[False, False, True], [False, False, True]]).to_sparse()
mt = masked_tensor(values, mask)

print("values:\n", values)
print("mask:\n", mask)
print("mt:\n", mt)

values:
 tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([3, 4, 5]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo)
mask:
 tensor(indices=tensor([[0, 1],
                       [2, 2]]),
       values=tensor([True, True]),
       size=(2, 3), nnz=2, layout=torch.sparse_coo)
mt:
 MaskedTensor(
  [
    [      --,       --, 3],
    [      --,       --, 5]
  ]
)

除了使用 torch.Tensor.to_sparse() 之外，我们还可以直接创建稀疏 COO 张量，这会给我们带来一个警告

警告

当使用像 MaskedTensor.to_sparse_coo()（类似于 Tensor.to_sparse()）这样的函数时，如果用户未像上述示例中那样指定索引，则默认情况下 0 值将为“未指定”。

下面，我们显式指定了 0 值

i = [[0, 1, 1],
     [2, 0, 2]]
v = [3, 4, 5]
m = torch.tensor([True, False, True])
values = torch.sparse_coo_tensor(i, v, (2, 3))
mask = torch.sparse_coo_tensor(i, m, (2, 3))
mt2 = masked_tensor(values, mask)

print("values:\n", values)
print("mask:\n", mask)
print("mt2:\n", mt2)

values:
 tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([3, 4, 5]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo)
mask:
 tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([ True, False,  True]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo)
mt2:
 MaskedTensor(
  [
    [      --,       --, 3],
    [      --,       --, 5]
  ]
)

请注意，mt 和 mt2 在表面上看起来相同，并且在绝大多数操作中，会产生相同的结果。但这给我们带来了一个关于实现的细节

data 和 mask – 仅对于稀疏 MaskedTensor – 可以在 创建时 具有不同数量的元素 (nnz())，但 mask 的索引必须是 data 索引的子集。在这种情况下，data 将通过 data = data.sparse_mask(mask) 采用 mask 的形状；换句话说，data 中任何在 mask 中不是 True 的元素（即未指定的元素）都将被丢弃。

因此，在底层，数据看起来略有不同；mt2 屏蔽掉了 “4” 值，而 mt 完全没有它。它们的底层数据具有不同的形状，这将使 mt + mt2 等操作无效。

print("mt data:\n", mt.get_data())
print("mt2 data:\n", mt2.get_data())

mt data:
 tensor(indices=tensor([[0, 1],
                       [2, 2]]),
       values=tensor([3, 5]),
       size=(2, 3), nnz=2, layout=torch.sparse_coo)
mt2 data:
 tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([3, 4, 5]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo)

稀疏 CSR 构造

我们还可以使用稀疏 CSR 张量构造稀疏 CSR MaskedTensor，并且像上面的示例一样，这会在底层产生类似的处理。

crow_indices = torch.tensor([0, 2, 4])
col_indices = torch.tensor([0, 1, 0, 1])
values = torch.tensor([1, 2, 3, 4])
mask_values = torch.tensor([True, False, False, True])

csr = torch.sparse_csr_tensor(crow_indices, col_indices, values, dtype=torch.double)
mask = torch.sparse_csr_tensor(crow_indices, col_indices, mask_values, dtype=torch.bool)
mt = masked_tensor(csr, mask)

print("mt:\n", mt)
print("mt data:\n", mt.get_data())

mt:
 MaskedTensor(
  [
    [  1.0000,       --],
    [      --,   4.0000]
  ]
)
mt data:
 tensor(crow_indices=tensor([0, 2, 4]),
       col_indices=tensor([0, 1, 0, 1]),
       values=tensor([1., 2., 3., 4.]), size=(2, 2), nnz=4,
       dtype=torch.float64, layout=torch.sparse_csr)

结论

在本教程中，我们介绍了如何将 MaskedTensor 与稀疏 COO 和 CSR 格式一起使用，并讨论了底层的一些细微之处，以防用户决定直接访问底层数据结构。稀疏存储格式和掩码语义确实具有很强的协同作用，以至于它们有时被用作彼此的代理（我们将在下一个教程中看到）。未来，我们肯定会计划继续投入并继续朝这个方向发展。

进一步阅读

要继续学习更多内容，您可以找到我们的 使用 MaskedTensor 为 Adagrad 高效编写“稀疏”语义教程，以查看 MaskedTensor 如何通过原生掩码语义简化现有工作流程的示例。