• 教程 >
  • (原型) 使用 MaskedTensor 高效编写 Adagrad 的“稀疏”语义
快捷方式

(原型) 使用 MaskedTensor 高效编写 Adagrad 的“稀疏”语义

在学习本教程之前,请先复习 MaskedTensor 概述稀疏性 教程。

引言和动机

问题 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)

稀疏的普通张量实现为

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_ 实现应该会导致致密化,但为了性能,我们强制它不要这样做。对于这种情况来说是可以的..直到我们想要引入新的压缩方案,例如 CSCBSRBSC。然后,我们将需要为每个方案引入单独的张量类型,并为使用不同存储格式压缩的梯度编写变体,这很不方便,而且扩展性和简洁性都不佳。

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)
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,       --]
  ]
)
print("param:\n", param)
print("param2:\n", param2)
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]])

结论

在本教程中,我们讨论了本机掩码语义如何能够为 PyTorch 中 Adagrad 的现有实现提供更简洁的开发体验,该实现使用稀疏性作为编写掩码语义的代理。但更重要的是,通过 MaskedTensor 允许掩码语义成为一等公民,消除了对稀疏性或不可靠的 hack 来模拟掩码的依赖,从而允许适当的独立性和开发,同时启用稀疏语义,例如此语义。

进一步阅读

要继续学习更多内容,您可以在 MaskedTensor 高级语义 中找到我们最终的回顾(目前为止),以了解 MaskedTensor 和 NumPy 的 MaskedArray 在设计决策方面的一些差异,以及约简语义。

脚本的总运行时间:(0 分 0.018 秒)

由 Sphinx-Gallery 生成的图库

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源