ONNX 操作符可微性标签简要指南

可微性标签

每个 ONNX 操作符的模式都包含一个针对每个输入和输出的可微性标签。本文档解释了这个标签的含义以及如何确保标签的正确性。简而言之,该标签标识了一个操作符的可微分输入和可微分输出的集合。标签的含义是,每个可微分输出对每个可微分输入的偏导数都是确定的。

定义可微性标签的方法

操作符的可微性定义包括几个方面。

  • 可微分输入,可以在 Gradient 的 xs 属性中引用。

  • 可微分输出,可以在 Gradient 的 y 属性中引用。

  • 用于计算 Jacobian 矩阵(或张量)的数学方程。一个变量(输入或输出)是否可微分由数学来判断。如果 Jacobian 矩阵(或张量)存在,那么所考虑的操作符就具有一些可微分输入和输出。

有几种实现自动微分的策略,例如前向累积、后向累积和对偶变量。由于大多数深度学习框架是基于后向的,审阅者应确保标签的 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:手动进行数学推导

第二种方法是正式证明从输出到输入的 Jacobian 矩阵(或张量)的存在性,并至少提供两个数值示例。在这种情况下,审阅者应仔细检查数学推导并确认数值结果是否正确。作者应添加足够的细节,以便任何 STEM 专业毕业的学生都能轻松审阅。

例如,为了展示 Add 的可微性,作者可以首先写下其方程:

C = A + B

为简单起见,假设 AB 是同形向量。

A = [a1, a2]^T
B = [b1, b2]^T
C = [c1, c2]^T

这里我们使用符号 ^T 表示所附矩阵或向量的转置。设 X = [a1, a2, b1, b2]^TY = [c1, c2]^T,并将 Add 视为一个将 X 映射到 Y 的函数。那么,该函数的 Jacobian 矩阵是一个 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]^TdL/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]^TdL/dB = [0.2, 0.8]^T。请注意,从 dL/dC 计算 dL/dAdL/dB 的过程通常称为操作符的后向(backward)。我们可以看到 Add 的后向操作符接受 dL/dC 作为输入,并生成两个输出 dL/dAdL/dB。因此,ABC 都可微分。通过将张量展平为一维向量,当不需要形状广播时,此示例可以扩展到所有张量。如果发生广播,广播元素的梯度是其**非广播**情况下所有相关元素的梯度之和。我们再次考虑上面的示例。如果 B = [b]^T 变成一个单元素向量,B 可能会广播到 [b1, b2]^T,那么 dL/dB = [dL/ db]^T = [dL/db1 + dL/db2]^T。对于高维张量,这实际上是沿着所有扩展轴进行的 ReduceSum 操作。