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维卷积运算符，带有可选的空洞（dilation）。
封装了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的字符串情况下，每个字符按位置标识使用字符 ‘N’ 表示
lhs、rhs和输出中的批次维度，使用字符 ‘C’ 表示 lhs 和输出中的特征维度，
使用字符 ‘I’ 和 ‘O’ 分别表示 rhs 中的输入和输出特征维度，以及
使用任何不同的字符表示 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 卷积）。