jax.lax.conv_general_dilated

jax.lax.conv_general_dilated#

jax.lax.conv_general_dilated(lhs, rhs, window_strides, padding, lhs_dilation=None, rhs_dilation=None, dimension_numbers=None, feature_group_count=1, batch_group_count=1, precision=None, preferred_element_type=None)[源代码][源代码]#

通用 n 维卷积操作符，可选的膨胀。

封装了 XLA 的 Conv 操作符。

参数:

  • lhs (Array) – 一个秩为 n+2 维的输入数组。

  • rhs (Array) – 一个秩为 n+2 的核权重数组。

  • window_strides (Sequence[int]) – 一个由 n 个整数组成的序列，表示窗口间的步幅。

  • padding (str | Sequence[tuple[int, int]]) – 可以是字符串 ‘SAME’、’SAME_LOWER’ 或 ‘VALID’，或者是一个包含 n 个 (low, high) 整数对的序列，这些对给出了在每个空间维度前后应用的填充。’SAME’ 和 ‘SAME_LOWER’ 添加填充以产生与输入相同大小的输出。填充在两侧之间平均分配或几乎平均分配。如果填充是奇数，额外的填充会添加在末尾用于 ‘SAME’，添加在开头用于 ‘SAME_LOWER’。

  • lhs_dilation (Sequence[int] | None) – None，或一个包含 n 个整数的序列，给出在 lhs 的每个空间维度上应用的扩张因子。LHS 扩张也称为转置卷积。

  • rhs_dilation (Sequence[int] | None) – None，或一个包含 n 个整数的序列，给出在 rhs 的每个空间维度上应用的膨胀因子。RHS 膨胀也称为空洞卷积。

  • dimension_numbers (ConvGeneralDilatedDimensionNumbers) – 可以是 None，一个 ConvDimensionNumbers 对象，或一个 3-元组 (lhs_spec, rhs_spec, out_spec)，其中每个元素是一个长度为 n+2 的字符串。

  • feature_group_count (int) – 整数，默认值为1。参见XLA HLO文档。

  • batch_group_count (int) – 整数，默认值为1。参见XLA HLO文档。

  • precision (lax.PrecisionLike) – 可选的。可以是 None，表示后端默认的精度，一个 Precision 枚举值（Precision.DEFAULT、Precision.HIGH 或 Precision.HIGHEST），一个字符串（例如 ‘highest’ 或 ‘fastest’，参见 jax.default_matmul_precision 上下文管理器），或者是一个包含两个 Precision 枚举值或字符串的元组，指示 lhs 和 rhs 的精度。

  • preferred_element_type (DTypeLike | None) – 可选。可以是 None，这意味着使用输入类型的默认累积类型，或者是一个数据类型，指示将结果累积到并返回该数据类型的结果。

返回:

包含卷积结果的数组。

返回类型:

Array

在 dimension_numbers 的字符串情况下，每个字符按位置标识：

  • lhs、rhs 和输出中的批处理维度用字符 ‘N’ 表示。

  • lhs 中的特征维度以及带有字符 ‘C’ 的输出

  • 在rhs中分别用字符’I’和’O’表示输入和输出的特征维度，并且

  • 使用任何不同字符表示 lhs、rhs 和输出之间的空间维度对应关系。下面的示例使用 ‘W’ 和 ‘H’。

例如，为了表示与具有两个空间维度的 conv 函数一致的维度编号，可以使用 ('NCHW', 'OIHW', 'NCHW')。再比如，为了表示与 TensorFlow Conv2D 操作一致的维度编号，可以使用 ('NHWC', 'HWIO', 'NHWC')。在使用后一种卷积维度规范形式时，窗口步幅与空间维度字符标签相关联，关联顺序按照这些标签在 rhs_spec 字符串中出现的顺序，因此 window_strides[0] 与 rhs_spec 中第一个不是 'I' 或 'O' 的字符对应的维度相匹配。

如果 dimension_numbers 是 None，默认值为 ``(‘NCHW’, ‘OIHW’, ‘NCHW’)``（用于2D卷积）。

目录