（原型）使用 MaskedTensor 高效编写“稀疏”语义的 Adagrad ¶

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

在学习本教程之前，请先回顾 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（功能）的实现，最终目标是简化并更忠实地表示掩码方法。

以下是没有掩码梯度或稀疏性的常规、密集代码路径的参考：

state_sum.addcmul_(grad, grad, value=1)
std = state_sum.sqrt().add_(eps)
param.addcdiv_(grad, std, value=-clr)

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

第三行倒数第二行 – 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 的实现更短更简单。特别是，围绕 _make_sparse （并且需要为每种布局实现一个单独的实现）的大部分样板代码都由 MaskedTensor 为用户处理。

到这一点，让我们打印这两个版本，以便更容易比较：

print("state_sum:\n", state_sum)
print("state_sum2:\n", state_sum2)

print("std:\n", std)
print("std2:\n", std2)

print("param:\n", param)
print("param2:\n", param2)

结论 ¶

在本教程中，我们讨论了原生掩码语义如何使 PyTorch 中 Adagrad 的现有实现拥有更清晰的开发者体验，该实现曾以稀疏性作为掩码语义的代理。但更重要的是，通过 MaskedTensor 将掩码语义提升为第一类公民，消除了对稀疏性或不可靠的技巧的依赖，从而实现了真正的独立和开发，同时使稀疏语义，如本例所示，成为可能。

进一步阅读

要继续学习更多，您可以查看我们关于 MaskedTensor 高级语义的最终评论（目前为止），以了解 MaskedTensor 和 NumPy 的 MaskedArray 在设计决策上的差异，以及缩减语义。

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

由 Sphinx-Gallery 生成的画廊