备注

点击此处下载完整示例代码

（原型）掩码张量稀疏性 ¶

创建时间：2025 年 4 月 1 日 | 最后更新时间：2025 年 4 月 1 日 | 最后验证：未验证

在开始本教程之前，请确保您已经阅读了我们的 MaskedTensor 概述教程。

简介

稀疏性在 PyTorch 中是一个快速发展和重要的领域；如果以下内容中的任何稀疏性术语令人困惑，请参阅稀疏性教程以获取更多详细信息。

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

备注

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

原则 ¶

在创建稀疏张量的 MaskedTensor 时，必须遵循一些原则：

data 和 mask 必须具有相同的存储格式，无论是 torch.strided 、 torch.sparse_coo 还是 torch.sparse_csr
data 和 mask 必须具有相同的大小，由 size() 指示

稀疏 COO 张量 ¶

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

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

indices : 大小为 (ndim, nse) 的数组，数据类型为 torch.int64
values : 大小为 (nse,) 的数组，可以是任何整数或浮点数数据类型

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

对于稀疏的 COO 和 CSR 张量，您可以通过以下方式构建一个 MaskedTensor ：

masked_tensor(sparse_tensor_data, sparse_tensor_mask)
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())

稀疏 CSR 张量

类似地， MaskedTensor 也支持 CSR（压缩稀疏行）稀疏张量格式。与稀疏 COO 张量存储索引元组不同，稀疏 CSR 张量旨在通过存储压缩行索引来降低内存需求。特别是，CSR 稀疏张量由三个一维张量组成：

crow_indices : 压缩行索引数组，大小为 (size[0] + 1,) 。该数组指示值数组中每个条目所在的行。最后一个元素是指定元素的数量，而 crow_indices[i+1] - crow_indices[i] 表示第 i 行中指定元素的数量。
col_indices : 大小为 (nnz,) 的数组。指示每个值的列索引。
values : 大小为 (nnz,) 的数组。包含 CSR 张量的值。

值得注意的是，稀疏 COO 和 CSR 张量都处于测试阶段。

例如：

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.sin()

二元

支持二进制运算符，但两个掩码张量的输入掩码必须匹配。关于为何做出此决定的更多信息，请参阅我们的 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)

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

减少

最后，支持了减少：

mt

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

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.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)

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())

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())

附录

稀疏 COO 构造

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

或者，我们也可以直接通过传入两个稀疏的 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)

我们可以使用 torch.Tensor.to_sparse() （类似于@1#）这样的函数，而不是直接创建稀疏的 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)

注意， mt 和 mt2 在表面上看起来是相同的，并且在绝大多数操作中，它们会产生相同的结果。但这让我们想到了实现中的一个细节：

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

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

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

稀疏 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())

结论 ¶

在本教程中，我们介绍了如何使用 MaskedTensor 与稀疏 COO 和 CSR 格式，并讨论了用户决定直接访问底层数据结构时的一些细节。稀疏存储格式和掩码语义确实具有很强的协同作用，以至于它们有时被用作对方的代理（我们将在下一教程中看到）。未来，我们当然计划在这个方向上进行投资和持续开发。

进一步阅读

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

脚本总运行时间：（0 分钟 0.000 秒）

由 Sphinx-Gallery 生成的画廊