ONNX 运算符可微性标签的简短指南¶

可微性标签¶

每个ONNX操作符的模式包括每个输入和输出的可微性标签。在本文档中，我们解释了此标签的含义以及如何确保标签的正确性。简而言之，该标签标识了操作符的可微输入和可微输出的集合。标签的含义是，每个可微输出的偏导数相对于每个可微输出是定义的。

定义可微性标签的方法¶

算子的可微性定义包含几个方面。

可微分的输入，可以在Gradient的xs属性中引用。
可微分的输出，可以在Gradient的y属性中引用。
计算雅可比矩阵（或张量）的数学方程。变量（输入或输出）是否可微由数学判断。如果雅可比矩阵（或张量）存在，则所考虑的算子具有一些可微的输入和输出。

实现自动微分有几种策略，如前向累积、后向累积和对偶变量。由于大多数深度学习框架都是基于后向的，评审者应确保标签的PR作者提供足够的细节。我们下面介绍几种方法来验证ONNX算子的可微性。

方法1：重用现有的深度学习框架¶

第一种方法是展示所考虑的运算符的反向操作存在于现有的框架中，如Pytorch或Tensorflow。在这种情况下，作者应提供一个可运行的python脚本，该脚本计算所考虑运算符的反向传递。作者还应指出如何将Pytorch或Tensor代码映射到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)

方法2：手动进行计算¶

第二种方法是正式证明从输出到输入的雅可比矩阵（或张量）的存在性，并至少提供两个数值示例。在这种情况下，审稿人应仔细检查数学部分，并确认数值结果是否正确。作者应添加足够的细节，以便任何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 的函数。那么，这个函数的雅可比矩阵是一个 4 行 2 列的矩阵，

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]]

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 操作。