unstable/maskedtensor_sparsity

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

MaskedTensor 稀疏性

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

简介

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

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

注意

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

原则

当创建带有稀疏张量的 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 附录。

# Disable prototype warnings and such

稀疏 CSR 张量

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

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

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

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

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

举例来说：

支持的操作

一元运算

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

二元运算

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

请看下面的例子

归约

最后，归约操作受支持

MaskedTensor 辅助方法

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

设置

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

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

MaskedTensor.is_sparse_coo()

MaskedTensor.is_sparse_csr()

附录

稀疏 COO 构造

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

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

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

警告

当使用像 MaskedTensor.to_sparse_coo()（类似于 Tensor.to_sparse()）这样的函数时，如果用户没有像上面示例那样指定索引，那么 0 值将默认“未指定”。

下面，我们显式指定 0

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

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

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

稀疏 CSR 构造

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

结论

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

进一步阅读

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

# %%%%%%RUNNABLE_CODE_REMOVED%%%%%%