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)[source]#

具有可选膨胀的通用 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) – None 或 n 个整数的序列，用于指定 lhs 每个空间维度的膨胀因子。LHS 膨胀也称为转置卷积。

  • rhs_dilation (Sequence[int] | None | None) – None 或 n 个整数的序列，用于指定 rhs 每个空间维度的膨胀因子。RHS 膨胀也称为空洞卷积。

  • dimension_numbers (ConvGeneralDilatedDimensionNumbers | None) – 可以是 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) – 可选。可以是 None（表示后端默认精度）、Precision 枚举值（Precision.DEFAULT、Precision.HIGH 或 Precision.HIGHEST）、字符串（例如 ‘highest’ 或 ‘fastest’，参见 jax.default_matmul_precision 上下文管理器），或两个 Precision 枚举或字符串的元组，分别表示 lhs 和 rhs 的精度。

  • preferred_element_type (DTypeLike | None | None) – 可选。可以是 None（表示输入类型的默认累加类型），或数据类型，表示将结果累加到该数据类型并返回该数据类型的结果。

返回:

包含卷积结果的数组。

返回类型:

数组

在 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 卷积）。

目录