注意
转到页面末尾下载完整的示例代码。
使用 MaskedTensor 高效编写 Adagrad 的“稀疏”语义#
在开始本教程之前,请回顾 MaskedTensor 的概述 (Overview) 和稀疏性 (Sparsity) 教程。
简介与动机#
Issue 1369 讨论了在编写 Adagrad 的“稀疏”语义时引入的额外代码行,但实际上,这些代码将稀疏性作为掩码语义的替代品,而非稀疏性的本意:即一种压缩和优化技术。此前,我们通过引入一次性语义和运算符来绕过缺乏正式掩码语义的问题,同时强制用户了解诸如索引和值之类的存储细节。
现在我们有了掩码语义,能更好地指出何时将稀疏性用作语义扩展。我们将把这一点与使用 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 的添加从技术上讲应该应用于所有值,但实际上只应用于指定的值。在这里,我们将稀疏性用作语义扩展,以强制执行某种定义的和未定义的值模式。如果梯度值的部分为零,即使它们可以通过其他稀疏存储布局进行压缩,如果被实例化,它们仍会被包含在内。这在理论上是非常脆弱的!话虽如此,有人可能会争辩说 eps 总是一个非常小的值,因此在实践中可能影响不大。
此外,将 add_ 实现为一种存储布局和压缩方案的稀疏性应该会导致稠密化(densification),但为了性能我们强制它不这样做。对于这种一次性情况,这没问题……直到我们想要引入新的压缩方案,例如 CSC、BSR 或 BSC。届时,我们需要为每种方案引入单独的 Tensor 类型,并为使用不同存储格式压缩的梯度编写变体,这既不方便,也不易扩展或简洁。
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 为用户处理了。
现在,让我们打印此版本和原始版本,以便进行更轻松的比较
结论#
在本教程中,我们讨论了原生掩码语义如何为 PyTorch 中现有的 Adagrad 实现提供更简洁的开发体验(它之前将稀疏性作为编写掩码语义的替代品)。更重要的是,通过 MaskedTensor 使掩码语义成为一等公民,消除了对稀疏性或不可靠的模拟掩码技巧的依赖,从而实现了适当的独立性和开发,同时启用了诸如本文所述的稀疏语义。
进一步阅读#
要继续深入了解,您可以查看我们关于 MaskedTensor 高级语义 (MaskedTensor Advanced Semantics) 的最终回顾(目前而言),以了解 MaskedTensor 与 NumPy 的 MaskedArray 在设计决策上的差异,以及归约语义。
# %%%%%%RUNNABLE_CODE_REMOVED%%%%%%
脚本总运行时间:(0 分 0.002 秒)