注意
跳转至页面底部下载完整示例代码。
使用 MaskedTensor 高效编写 Adagrad 的“稀疏”语义#
在学习本教程之前，请先查阅 MaskedTensor 概述 和 稀疏性 (Sparsity) 教程。
介绍与动机#
Issue 1369 讨论了在编写 Adagrad 的“稀疏”语义时引入的额外代码行。但实际上，该代码使用稀疏性作为掩码语义的替代品，而非稀疏性的本意：即一种压缩和优化技术。此前，我们通过引入一次性 (one-off) 的语义和算子来解决缺乏正式掩码语义的问题，同时被迫让用户了解诸如索引和值等存储细节。
现在我们有了掩码语义，能更好地指出何时将稀疏性用作语义扩展。我们还将通过对比，将其与使用 MaskedTensor 编写的等效代码进行比较。最后，我们将重复这些代码片段且不加额外注释，以展示其在简洁性上的差异。
准备工作#
# Disable prototype warnings and such
# Some hyperparameters
使用 MaskedTensor 简化代码#
在深入细节之前，让我们更具体地介绍一下问题。我们将深入研究 PyTorch 中的 Adagrad (函数式) 实现，其最终目标是简化并更忠实地呈现掩码方法。
作为参考，这是没有掩码梯度或稀疏性的常规稠密代码路径
state_sum.addcmul_(grad, grad, value=1)
std = state_sum.sqrt().add_(eps)
param.addcdiv_(grad, std, value=-clr)
稀疏张量的原生实现是
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) 它避免了在稠密张量和稀疏张量之间进行转换。
原始稀疏实现#
现在，让我们通过一些行内注释来解析代码
# We don't support sparse gradients
# pow(2) has the same semantics for both sparse and dense memory layouts since 0^2 is zero
# 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`.
# 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!
倒数第三行 – std = state_sum.sparse_mask(grad) – 是我们出现重大分歧的地方。
eps 的添加从技术上讲应该应用于所有值，但实际上仅应用于指定的值。在这里，我们将稀疏性用作语义扩展，以强制执行某种定义值和未定义值的模式。如果梯度值的某些部分为零，即使它们可以通过其他稀疏存储布局进行压缩，如果被物化 (materialized)，它们仍会被包含在内。这在理论上是非常脆弱的！话虽如此，有人可能会争辩说 eps 总是很小，因此在实践中可能并不重要。
此外，针对作为存储布局和压缩方案的稀疏性实现 add_ 本应导致稠密化，但为了性能，我们强制其不这样做。对于这种一次性情况，这样做是可以的……直到我们想要引入新的压缩方案，例如 CSC、BSR 或 BSC。届时，我们需要为每种方案引入单独的张量类型，并为使用不同存储格式压缩的梯度编写变体，这既不方便，也不具备可扩展性或清晰度。
MaskedTensor 稀疏实现#
我们一直将作为优化的稀疏性与作为 PyTorch 语义扩展的稀疏性混为一谈。MaskedTensor 建议将稀疏优化与语义扩展解耦；例如，目前我们无法拥有具备稀疏存储的稠密语义，或具备稠密存储的掩码语义。MaskedTensor 通过刻意分离存储与语义，实现了这些想法。
考虑使用掩码梯度的上述示例
# Let's now import MaskedTensor!
# Create an entirely new set of parameters to avoid errors
# 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
请注意，这些实现看起来非常相似，但 MaskedTensor 的实现更短、更简洁。特别是，围绕 _make_sparse 的大部分样板代码（以及需要为每种布局编写单独实现的需要）都由 MaskedTensor 为用户处理了。
现在，让我们打印此版本和原始版本以便进行比较
结论#
在本教程中，我们讨论了原生掩码语义如何为 Adagrad 在 PyTorch 中的现有实现提供更清晰的开发体验，该实现此前将稀疏性用作编写掩码语义的代理。但更重要的是，允许掩码语义通过 MaskedTensor 成为“一等公民”，消除了对稀疏性或不可靠 hack 手段的依赖，从而实现了真正的独立性和开发性，同时启用了诸如本文所述的稀疏语义。
进一步阅读#
要继续深入了解，您可以查看我们（目前为止）的最终回顾：MaskedTensor 高级语义，以了解 MaskedTensor 与 NumPy 的 MaskedArray 在设计决策上的一些差异，以及归约 (reduction) 语义。
# %%%%%%RUNNABLE_CODE_REMOVED%%%%%%
脚本总运行时间：（0 分 0.002 秒）