ONNX 算子可微分标签简明指南¶
可微分标签¶
每个 ONNX 算子的模式(schema)都包含一个关于每个输入和输出的可微分标签。本文档将解释此标签的含义以及如何确保标签的正确性。简而言之,该标签标识了一个算子的可微分输入和可微分输出的集合。该标签的含义是:每个可微分输出的偏导数相对于每个可微分输出是存在的。
定义可微分标签的方式¶
算子的可微分定义包含几个方面。
可微分输入,可以在 Gradient 的
xs
属性中引用。可微分输出,可以在 Gradient 的
y
属性中引用。计算雅可比矩阵(或张量)的数学方程。一个变量(输入或输出)是否可微分是通过数学来判断的。如果雅可比矩阵(或张量)存在,那么该算子就存在一些可微分的输入和输出。
实现自动微分有几种策略,如前向累积、后向累积和双变量法。由于大多数深度学习框架都基于后向传播,评审者应确保 PR 作者在标签方面提供足够详细的信息。下面我们将介绍两种验证 ONNX 算子可微分性的方法。
方法一:复用现有的深度学习框架¶
第一种方法是展示所考虑的算子在现有框架(如 Pytorch 或 Tensorflow)中的后向(backward)操作是否存在。在这种情况下,作者应提供一个可运行的 Python 脚本,该脚本计算所考虑算子的后向传播。作者还应指出如何将 Pytorch 或 Tensorflow 代码映射到 ONNX 格式(例如,作者可以调用 torch.onnx.export
来保存 ONNX 模型)。下面的脚本展示了使用 Pytorch 对 ONNX Reshape 进行可微分性验证。
import torch
import torch.nn as nn
# A single-operator model. It's literally a Pytorch Reshape.
# Note that Pytorch Reshape can be directly mapped to ONNX Reshape.
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
def forward(self, x):
y = torch.reshape(x, (x.numel(),))
y.retain_grad()
return y
model = MyModel()
x = torch.tensor([[1., -1.], [1., 1.]], requires_grad=True)
y = model(x)
dy = torch.tensor([1., 2., 3., 4.])
torch.autograd.backward([y],
grad_tensors=[dy],
retain_graph=True,
create_graph=True,
grad_variables=None)
# This example shows the input and the output in Pytorch are differentiable.
# From the exported ONNX model below, we also see that "x" is the first input
# of ONNX Reshape and "y" the output of ONNX Reshape. Therefore, we can say
# the first input and the output of ONNX Reshape are differentiable.
print(x.grad)
print(y.grad)
with open('model.onnx', 'wb') as f:
torch.onnx.export(model, x, f)
方法二:手动进行数学推导¶
第二种方法是通过至少两个数值示例,正式证明从输出到输入的雅可比矩阵(或张量)的存在性。在这种情况下,评审者应仔细审查数学推导,并确认数值结果的正确性。作者应提供足够的细节,以便任何 STEM 专业的毕业生都能轻松理解。
例如,为了展示 Add 算子的可微分性,作者可以首先写下其方程
C = A + B
为简单起见,假设 A
和 B
是相同形状的向量。
A = [a1, a2]^T
B = [b1, b2]^T
C = [c1, c2]^T
这里我们使用符号 ^T
来表示其附加矩阵或向量的转置。设 X = [a1, a2, b1, b2]^T
和 Y = [c1, c2]^T
,并将 Add 视为一个将 X
映射到 Y
的函数。那么,该函数的雅可比矩阵是一个 4x2 的矩阵:
J = [[dc1/da1, dc2/da1],
[dc1/da2, dc2/da2],
[dc1/db1, dc2/db1],
[dc1/db2, dc2/db2]]
= [[1, 0],
[0, 1],
[1, 0],
[0, 1]]
If
dL/dC = [dL/dc1, dL/dc2]^T,
则 dL/dA = [dL/da1, dL/da2]^T
和 dL/dB = [dL/db1, dL/db2]^T
可以从以下元素计算得出:
[[dL/da1], [dL/da2], [dL/db1], [dL/db2]]
= J * dL/dC
= [[dL/dc1], [dL/dc2], [dL/dc1], [dL/dc2]]
其中 *
表示标准的矩阵乘法。如果 dL/dC = [0.2, 0.8]^T
,那么 dL/dA = [0.2, 0.8]^T
和 dL/dB = [0.2, 0.8]^T
。请注意,从 dL/dC
计算 dL/dA
和 dL/dB
的过程通常被称为算子的后向传播。我们可以看到,Add 算子的后向操作以 dL/dC
作为输入,并产生两个输出 dL/dA
和 dL/dB
。因此,A
、B
和 C
都是可微分的。通过将张量展平成一维向量,这个例子可以扩展到覆盖所有张量,前提是没有形状广播。如果发生广播,广播元素的梯度是其在 **非广播** 情况下的所有相关元素的梯度之和。让我们再次考虑上面的例子。如果 B = [b]^T
变成一个单元素向量,B
可能会被广播到 [b1, b2]^T
,并且 dL/dB = [dL/ db]^T = [dL/db1 + dL/db2]^T
。对于高维张量,这实际上是沿着所有扩展轴的 ReduceSum 操作。