jax.scipy.signal.convolve2d

内容

jax.scipy.signal.convolve2d#

jax.scipy.signal.convolve2d(in1, in2, mode='full', boundary='fill', fillvalue=0, precision=None)[source]#

两个二维数组的卷积。

JAX 实现的 scipy.signal.convolve2d().

参数:
  • in1 (Array) – 卷积的左侧输入。必须满足 in1.ndim == 2

  • in2 (Array) – 卷积的右侧输入。必须具有 in2.ndim == 2

  • mode (str) –

    控制输出的大小。可用操作包括

    • "full": (默认) 输出输入的完整卷积。

    • "same": 返回 "full" 输出的中心部分,其大小与 in1 相同。

    • "valid": 返回 "full" 输出中不依赖于数组边缘填充的部分。

  • boundary (str) – 仅支持 "fill"

  • fillvalue (float) – 仅支持 0

  • method

    控制计算方法。选项包括

    • "auto": (默认) 始终使用 "direct" 方法。

    • "direct": 降级到 jax.lax.conv_general_dilated().

    • "fft": 通过快速傅里叶变换计算结果。

  • precision (PrecisionLike | None) – 指定计算的精度。有关可用值的说明,请参阅 jax.lax.Precision.

返回值:

包含卷积结果的数组。

返回类型:

Array

另请参阅

示例

一些 2D 卷积示例

>>> x = jnp.array([[1, 2],
...                [3, 4]])
>>> y = jnp.array([[2, 1, 1],
...                [4, 3, 4],
...                [1, 3, 2]])

完整 2D 卷积使用边缘处的隐式零填充

>>> jax.scipy.signal.convolve2d(x, y, mode='full')
Array([[ 2.,  5.,  3.,  2.],
       [10., 22., 17., 12.],
       [13., 30., 32., 20.],
       [ 3., 13., 18.,  8.]], dtype=float32)

指定 mode = 'same' 返回与第一个输入大小相同的中心 2D 卷积

>>> jax.scipy.signal.convolve2d(x, y, mode='same')
Array([[22., 17.],
       [30., 32.]], dtype=float32)

指定 mode = 'valid' 仅返回两个数组完全重叠的 2D 卷积部分

>>> jax.scipy.signal.convolve2d(x, y, mode='valid')
Array([[22., 17.],
       [30., 32.]], dtype=float32)