(原型) MaskedTensor 稀疏性

Created On: Oct 28, 2022 | Last Updated: Dec 12, 2023 | Last Verified: Not Verified

开始学习本教程之前，请确保已查看我们的`MaskedTensor概述教程 <https://pytorch.org/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()

第二种方法更容易说明，因此我们在下面展示了该方法，但有关第一种方法和方法细节的更多信息，请阅读 Sparse 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 张量均处于 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.sin()

二元运算

二元运算符 也受支持，但两个 masked 张量的输入屏蔽必须匹配。有关为什么做出此决定的更多信息，请参阅我们的 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() – 检查 :class:`MaskedTensor`&apos; 的布局是否匹配任何支持的稀疏布局（目前为 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，然后通过 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)

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

警告

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

下面，我们明确指定了 0&apos;：

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 = data.sparse_mask(mask)，data 将采用 mask 的形状；换句话说，mask 中不是 True 的元素（即未指定的元素）将在 data 中被丢弃。

因此，从内部来看，数据看起来稍有不同；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())

总结

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

进一步阅读

要继续学习，您可以查看我们的 Efficiently writing “sparse” semantics for Adagrad with MaskedTensor tutorial ，以了解如何使用 MaskedTensor 将现有工作流简化为原生掩码语义的示例。