注意
点击此处下载完整的示例代码
（原型）使用 MaskedTensor 为 Adagrad 高效编写“稀疏”语义¶
创建于：2022 年 10 月 28 日 | 最后更新：2022 年 10 月 28 日 | 最后验证：未验证
在学习本教程之前，请回顾 MaskedTensor 概述 和 稀疏性 教程。
简介和动机¶
Issue 1369 讨论了在为 Adagrad 编写“稀疏”语义时引入的额外代码行，但实际上，该代码使用稀疏性作为掩码语义的代理，而不是稀疏性的预期用例：一种压缩和优化技术。 以前，我们通过引入一次性语义和运算符来解决缺少正式掩码语义的问题，同时迫使用户了解存储细节，例如索引和值。
现在我们有了掩码语义，我们可以更好地指出何时将稀疏性用作语义扩展。 我们还将使用 MaskedTensor 编写的等效代码进行比较和对比。 最后，重复代码片段，不添加额外注释，以显示简洁性的差异。
准备工作¶
import torch
import warnings
# Disable prototype warnings and such
warnings.filterwarnings(action='ignore', category=UserWarning)
# Some hyperparameters
eps = 1e-10
clr = 0.1
i = torch.tensor([[0, 1, 1], [2, 0, 2]])
v = torch.tensor([3, 4, 5], dtype=torch.float32)
grad = torch.sparse_coo_tensor(i, v, [2, 4])
使用 MaskedTensor 简化代码¶
在我们深入细节之前，让我们更具体地介绍一下问题。我们将研究 PyTorch 中 Adagrad（函数式） 的实现，最终目标是简化并更忠实地表示掩码方法。
作为参考，这是没有掩码梯度或稀疏性的常规密集代码路径
稀疏的 vanilla 张量实现是
def _make_sparse(grad, grad_indices, values):
size = grad.size()
if grad_indices.numel() == 0 or values.numel() == 0:
return torch.empty_like(grad)
return torch.sparse_coo_tensor(grad_indices, values, size)
grad = grad.coalesce() # the update is non-linear so indices must be unique
grad_indices = grad._indices()
grad_values = grad._values()
state_sum.add_(_make_sparse(grad, grad_indices, grad_values.pow(2))) # a different _make_sparse per layout
std = state_sum.sparse_mask(grad)
std_values = std._values().sqrt_().add_(eps)
param.add_(_make_sparse(grad, grad_indices, grad_values / std_values), alpha=-clr)
而 MaskedTensor 将代码最小化到以下代码段
state_sum2 = state_sum2 + masked_grad.pow(2).get_data()
std2 = masked_tensor(state_sum2.to_sparse(), mask)
std2 = std2.sqrt().add(eps)
param2 = param2.add((masked_grad / std2).get_data(), alpha=-clr)
在本教程中，我们将逐行浏览每个实现，但乍一看，我们可以注意到 (1) MaskedTensor 实现的简洁程度，以及 (2) 它如何避免密集张量和稀疏张量之间的转换。
原始稀疏实现¶
现在，让我们分解代码并添加一些内联注释
def _make_sparse(grad, grad_indices, values):
size = grad.size()
if grad_indices.numel() == 0 or values.numel() == 0:
return torch.empty_like(grad)
return torch.sparse_coo_tensor(grad_indices, values, size)
# We don't support sparse gradients
param = torch.arange(8).reshape(2, 4).float()
state_sum = torch.full_like(param, 0.5) # initial value for state sum
grad = grad.coalesce() # the update is non-linear so indices must be unique
grad_indices = grad._indices()
grad_values = grad._values()
# pow(2) has the same semantics for both sparse and dense memory layouts since 0^2 is zero
state_sum.add_(_make_sparse(grad, grad_indices, grad_values.pow(2)))
# We take care to make std sparse, even though state_sum clearly is not.
# This means that we're only applying the gradient to parts of the state_sum
# for which it is specified. This further drives the point home that the passed gradient is not sparse, but masked.
# We currently dodge all these concerns using the private method `_values`.
std = state_sum.sparse_mask(grad)
std_values = std._values().sqrt_().add_(eps)
# Note here that we currently don't support div for sparse Tensors because zero / zero is not well defined,
# so we're forced to perform `grad_values / std_values` outside the sparse semantic and then convert back to a
# sparse tensor with `make_sparse`.
# We'll later see that MaskedTensor will actually handle these operations for us as well as properly denote
# undefined / undefined = undefined!
param.add_(_make_sparse(grad, grad_indices, grad_values / std_values), alpha=-clr)
tensor([[0.0000, 1.0000, 1.9027, 3.0000],
[3.9015, 5.0000, 5.9010, 7.0000]])
倒数第三行 – std = state_sum.sparse_mask(grad) – 是我们发生非常重要分歧的地方。
eps 的添加在技术上应应用于所有值，但实际上仅应用于指定的值。 在这里，我们使用稀疏性作为语义扩展，并强制执行已定义值和未定义值的特定模式。 如果梯度值的部分为零，即使它们可以被其他稀疏存储布局压缩，它们仍然会被包含在内（如果已物化）。 这在理论上非常脆弱！ 也就是说，有人可能会争辩说 eps 总是非常小，因此在实践中可能无关紧要。
此外，作为存储布局和压缩方案的稀疏性的 add_ 实现应该导致密集化，但为了性能，我们强制它不这样做。 对于这种一次性情况，这很好……直到我们想引入新的压缩方案，例如 CSC、BSR 或 BSC。 然后，我们将需要为每种方案引入单独的张量类型，并为使用不同存储格式压缩的梯度编写变体，这既不方便，也不可扩展或不简洁。
MaskedTensor 稀疏实现¶
我们一直将稀疏性作为优化与稀疏性作为 PyTorch 的语义扩展混为一谈。 MaskedTensor 建议将稀疏性优化与语义扩展分离； 例如，目前我们不能拥有具有稀疏存储的密集语义或具有密集存储的掩码语义。 MaskedTensor 通过有目的地将存储与语义分离来实现这些想法。
考虑上面使用掩码梯度的示例
# Let's now import MaskedTensor!
from torch.masked import masked_tensor
# Create an entirely new set of parameters to avoid errors
param2 = torch.arange(8).reshape(2, 4).float()
state_sum2 = torch.full_like(param, 0.5) # initial value for state sum
mask = (grad.to_dense() != 0).to_sparse()
masked_grad = masked_tensor(grad, mask)
state_sum2 = state_sum2 + masked_grad.pow(2).get_data()
std2 = masked_tensor(state_sum2.to_sparse(), mask)
# We can add support for in-place operations later. Notice how this doesn't
# need to access any storage internals and is in general a lot shorter
std2 = std2.sqrt().add(eps)
param2 = param2.add((masked_grad / std2).get_data(), alpha=-clr)
请注意，这些实现看起来非常相似，但 MaskedTensor 实现更短更简单。 特别是，用户使用 MaskedTensor 处理了围绕 _make_sparse（以及需要为每个布局单独实现）的大量样板代码。
此时，让我们打印此版本和原始版本，以便更轻松地进行比较
print("state_sum:\n", state_sum)
print("state_sum2:\n", state_sum2)
state_sum:
tensor([[ 0.5000, 0.5000, 9.5000, 0.5000],
[16.5000, 0.5000, 25.5000, 0.5000]])
state_sum2:
tensor([[ 0.5000, 0.5000, 9.5000, 0.5000],
[16.5000, 0.5000, 25.5000, 0.5000]])
print("std:\n", std)
print("std2:\n", std2)
std:
tensor(indices=tensor([[0, 1, 1],
[2, 0, 2]]),
values=tensor([3.0822, 4.0620, 5.0498]),
size=(2, 4), nnz=3, layout=torch.sparse_coo)
std2:
MaskedTensor(
[
[ --, --, 3.0822, --],
[ 4.0620, --, 5.0498, --]
]
)
param:
tensor([[0.0000, 1.0000, 1.9027, 3.0000],
[3.9015, 5.0000, 5.9010, 7.0000]])
param2:
tensor([[0.0000, 1.0000, 1.9027, 3.0000],
[3.9015, 5.0000, 5.9010, 7.0000]])
结论¶
在本教程中，我们讨论了原生掩码语义如何为 Adagrad 在 PyTorch 中的现有实现提供更简洁的开发者体验，该实现使用稀疏性作为编写掩码语义的代理。 但更重要的是，允许掩码语义通过 MaskedTensor 成为一等公民，消除了对稀疏性或不可靠的 hack 来模仿掩码的依赖，从而实现了适当的独立性和开发，同时实现了稀疏语义，例如这种语义。
进一步阅读¶
要继续了解更多信息，您可以找到我们（目前）关于 MaskedTensor 高级语义 的最终评论，以了解 MaskedTensor 和 NumPy 的 MaskedArray 之间在设计决策以及归约语义方面的一些差异。
脚本的总运行时间： (0 分 0.017 秒)
