ONNX 中的广播

在 ONNX 中,逐元素(element-wise)算子可以接受不同形状的输入,只要输入张量可以广播到相同的形状即可。ONNX 支持两种类型的广播:多向广播(multidirectional broadcasting)和单向广播(unidirectional broadcasting)。我们将在接下来的章节中分别介绍这两种广播类型。

多向广播

在 ONNX 中,如果满足以下任一条件,则一组张量可以进行多向广播到相同的形状

  • 所有张量具有完全相同的形状。

  • 所有张量具有相同的维度数,并且每个维度的长度要么是共同长度,要么是 1。

  • 维度数太少的张量可以在其形状前添加一个长度为 1 的维度以满足条件 2。

例如,以下张量形状支持多向广播

  • shape(A) = (2, 3, 4, 5), shape(B) = (,), 即 B 是一个标量 ==> shape(result) = (2, 3, 4, 5)

  • shape(A) = (2, 3, 4, 5), shape(B) = (5,), ==> shape(result) = (2, 3, 4, 5)

  • shape(A) = (4, 5), shape(B) = (2, 3, 4, 5), ==> shape(result) = (2, 3, 4, 5)

  • shape(A) = (1, 4, 5), shape(B) = (2, 3, 1, 1), ==> shape(result) = (2, 3, 4, 5)

  • shape(A) = (3, 4, 5), shape(B) = (2, 1, 1, 1), ==> shape(result) = (2, 3, 4, 5)

多向广播与 Numpy 的广播规则相同。

ONNX 中的以下算子支持多向广播

单向广播

在 ONNX 中,如果满足以下任一条件,则张量 B 可以单向广播到张量 A

  • 张量 A 和 B 具有完全相同的形状。

  • 张量 A 和 B 都具有相同的维度数,并且每个维度的长度要么是共同长度,要么是 B 的长度为 1。

  • 张量 B 的维度数太少,并且可以在其形状前添加一个长度为 1 的维度以满足条件 2。

当发生单向广播时,输出的形状与张量 A 的形状相同(即,两个输入张量中较大的形状)。

在以下示例中,张量 B 可以单向广播到张量 A

  • shape(A) = (2, 3, 4, 5), shape(B) = (,), 即 B 是一个标量 ==> shape(result) = (2, 3, 4, 5)

  • shape(A) = (2, 3, 4, 5), shape(B) = (5,), ==> shape(result) = (2, 3, 4, 5)

  • shape(A) = (2, 3, 4, 5), shape(B) = (2, 1, 1, 5), ==> shape(result) = (2, 3, 4, 5)

  • shape(A) = (2, 3, 4, 5), shape(B) = (1, 3, 1, 5), ==> shape(result) = (2, 3, 4, 5)

ONNX 中的以下算子支持单向广播