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 格式代表“坐标格式”（coordinate format），其中指定元素存储为索引元组及其对应的值。也就是说，提供以下内容：

• indices：尺寸为 (ndim, nse) 且数据类型为 torch.int64 的数组

• values：尺寸为 (nse,) 且具有任意整数或浮点数据类型的数组

其中 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（压缩稀疏行） 稀疏张量格式。与存储索引元组的稀疏 COO 张量不同，稀疏 CSR 张量旨在通过存储压缩行索引来减少内存需求。具体而言，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 的输入掩码（masks）必须匹配。有关做出此决定的详细信息，请查看我们的 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 值默认为“未指定”（unspecified）。

下面我们显式指定 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%%%%%%