注意
跳到末尾 下载完整的示例代码。
使用 MaskedTensor 高效编写 Adagrad 的“稀疏”语义#
在完成本教程之前,请先查阅 MaskedTensor 概述 和 稀疏性 教程。
引言和动机#
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_ 应该导致致密化,但为了性能我们强制它不这样做。对于这个一次性案例来说没问题……直到我们想引入新的压缩方案,例如 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
为用户处理了。
此时,让我们打印此版本和原始版本以便于比较
结论#
在本教程中,我们讨论了原生掩码语义如何为 PyTorch 中 Adagrad 的现有实现提供更清晰的开发体验,该实现使用稀疏性作为编写掩码语义的代理。但更重要的是,通过 MaskedTensor 允许掩码语义成为一等公民,消除了对稀疏性或不可靠的黑客手段来模拟掩码的依赖,从而实现了适当的独立性和开发,同时启用了稀疏语义,例如这种语义。
进一步阅读#
要继续学习更多内容,您可以在 MaskedTensor 高级语义 中找到我们(目前)的最终评论,以查看 MaskedTensor
和 NumPy 的 MaskedArray 之间的一些设计决策差异,以及约简语义。
# %%%%%%RUNNABLE_CODE_REMOVED%%%%%%
脚本总运行时间:(0 分 0.002 秒)