JAX 卷积完全指南：从 jnp.convolve 到 lax.conv_general_dilated 的 N 维通用卷积实战

JAX 卷积完全指南：从 jnp.convolve 到 lax.conv_general_dilated 的 N 维通用卷积实战

【免费下载链接】jaxComposable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/gh_mirrors/jax/jax

JAX 为卷积运算提供了从一维信号平滑到神经网络批量化多维卷积的完整接口栈：jax.numpy.convolve面向一维基础卷积，jax.scipy.signal.convolve/convolve2d覆盖 N 维与二维卷积，而jax.lax.conv_general_dilated则是可组合 stride、padding、dilation、维度重排与分组卷积的通用底层原语。本文将逐个层级讲解其用法与取舍，并深入源码剖析 dimension numbers、转置卷积（Transposed Convolution）与空洞卷积（Atrous Convolution）的实现原理，读完即可在自己的模型中自由切换 NHWC/NCHW、HWIO/OIHW 等任意数据布局。

接口总览：JAX 提供的四类卷积入口

JAX 计算卷积的接口分布在三个命名空间中，按通用程度递增：

• jax.numpy.convolve（以及jax.numpy.correlate）：一维卷积，接口对齐 NumPy；
• jax.scipy.signal.convolve（以及jax.scipy.signal.correlate）：N 维卷积，接口对齐 SciPy；
• jax.scipy.signal.convolve2d（以及jax.scipy.signal.correlate2d）：二维卷积；
• jax.lax.conv_general_dilated：最通用的批量 N 维卷积原语，构建深度神经网络时从这里起步。

对于基础的卷积操作，jax.numpy与jax.scipy的接口通常已经够用；当需要更一般的批量多维卷积（stride、dilation、分组、自定义维度布局）时，应转向jax.lax层的conv_general_dilated——它是下面所有便捷接口的底层实现。

从源码看，jax.scipy.signal.convolve的method='direct'（默认的'auto'也等同 direct）最终就是调用lax.conv_general_dilated完成计算（见 jax/_src/scipy/signal.py）：

result = lax.conv_general_dilated(in1[None, None], in2[None, None], strides, padding, precision=precision) return result[0, 0]

也就是说，所有高层接口最终都汇入同一个底层原语，掌握conv_general_dilated就等于掌握了 JAX 卷积的"发动机"。

一维卷积：jax.numpy.convolve

jax.numpy.convolve是 NumPynumpy.convolve的 JAX 接口（实现位于 jax/_src/numpy/lax_numpy.py），可以直接接入jit、vmap、grad等 JAX 变换。下面是一个典型的 1D 滑动平均平滑示例：

import matplotlib.pyplot as plt from jax import random import jax.numpy as jnp import numpy as np key = random.key(1701) x = jnp.linspace(0, 10, 500) y = jnp.sin(x) + 0.2 * random.normal(key, shape=(500,)) window = jnp.ones(10) / 10 y_smooth = jnp.convolve(y, window, mode='same') plt.plot(x, y, 'lightgray') plt.plot(x, y_smooth, 'black');

这里window = jnp.ones(10) / 10是一个长度为 10 的归一化矩形窗，卷积后得到平滑曲线。mode参数控制边界条件的处理方式：

• 'full'：输出完整卷积结果，长度为len(a) + len(v) - 1；
• 'same'：输出与第一个输入a等长（本例即 500），边界按零填充处理；
• 'valid'：只返回两个数组完全重叠的部分，输出长度更短。

示例中使用mode='same'保证输出与输入等长，从而可以直接与原信号y对齐绘图。更完整的边界选项可查阅jax.numpy.convolve的文档字符串，其语义与原始numpy.convolve完全一致。与之配对的jax.numpy.correlate实现的是互相关（不翻转核），与卷积互为"镜像"关系。

N 维卷积：jax.scipy.signal.convolve 与 convolve2d

jax.scipy.signal.convolve把一维卷积推广到任意 N 维，接口与jax.numpy.convolve类似。下面的例子用高斯核平滑含噪图像，实现去噪：

from scipy import misc import jax.scipy as jsp fig, ax = plt.subplots(1, 3, figsize=(12, 5)) # Load a sample image; compute mean() to convert from RGB to grayscale. image = jnp.array(misc.face().mean(-1)) ax[0].imshow(image, cmap='binary_r') ax[0].set_title('original') # Create a noisy version by adding random Gaussian noise key = random.key(1701) noisy_image = image + 50 * random.normal(key, image.shape) ax[1].imshow(noisy_image, cmap='binary_r') ax[1].set_title('noisy') # Smooth the noisy image with a 2D Gaussian smoothing kernel. x = jnp.linspace(-3, 3, 7) window = jsp.stats.norm.pdf(x) * jsp.stats.norm.pdf(x[:, None]) smooth_image = jsp.signal.convolve(noisy_image, window, mode='same') ax[2].imshow(smooth_image, cmap='binary_r') ax[2].set_title('smoothed');

核心操作只有一行：jsp.signal.convolve(noisy_image, window, mode='same')。其中高斯核由两个一维高斯 PDF 的外积jsp.stats.norm.pdf(x) * jsp.stats.norm.pdf(x[:, None])构造（7×7），mode='same'保证输出与输入图像同尺寸。

与一维情况一样，mode仍支持'full'/'same'/'valid'三种取值（默认'full'）。此外jax.scipy.signal.convolve还提供了method参数（见 jax/_src/scipy/signal.py）：

• method='direct'：直接下放到lax.conv_general_dilated计算，适合中小尺寸核；
• method='fft'：通过快速傅里叶变换（fftconvolve）计算，适合大核大输入，利用 FFT 将卷积转为频域乘法；
• method='auto'：当前实现始终走direct路径。

对于二维场景，还有专门的jax.scipy.signal.convolve2d（以及correlate2d），其mode语义同上，boundary目前仅支持'fill'、fillvalue仅支持0。

通用卷积：lax.conv_general_dilated

构建神经网络时常用的批量多维卷积远比基础卷积复杂：需要同时控制 batch 维度、多输入/输出通道、stride、padding、dilation 以及任意维度排列。XLA 提供了非常通用的 N 维conv_general_dilated算子，但它的参数直观性较差，本小节用一组完整示例讲透常见用法。关于卷积算子的家族综述，可参考卷积算术指南（A guide to convolutional arithmetic，arXiv:1603.07285）。

首先定义一个简单的对角边缘检测核：

# 2D kernel - HWIO layout kernel = jnp.zeros((3, 3, 3, 3), dtype=jnp.float32) kernel += jnp.array([[1, 1, 0], [1, 0,-1], [0,-1,-1]])[:, :, jnp.newaxis, jnp.newaxis] print("Edge Conv kernel:") plt.imshow(kernel[:, :, 0, 0]);

再构造一张合成的彩色图像（三个通道中各放一个方块）：

# NHWC layout img = jnp.zeros((1, 200, 198, 3), dtype=jnp.float32) for k in range(3): x = 30 + 60*k y = 20 + 60*k img = img.at[0, x:x+10, y:y+10, k].set(1.0) print("Original Image:") plt.imshow(img[0]);

注意这里刻意选择了NHWC图像布局与HWIO核布局，以展示conv_general_dilated在维度排列上的灵活性（默认约定其实更接近 NCHW/OIHW，见下文）。

便捷函数 lax.conv 与 lax.conv_with_general_padding

lax.conv和lax.conv_with_general_padding是conv_general_dilated的便捷封装（源码见 jax/_src/lax/convolution.py），它们的参数更少、上手更快，但有一个重要前提：

⚠️lax.conv与lax.conv_with_general_padding假设输入是NCHW图像、核是OIHW布局（lax.conv固定使用'SAME'/'VALID'字符串 padding；conv_with_general_padding额外接受显式的(low, high)padding 序列以及 lhs/rhs dilation）。

因此，需要先把上面的 NHWC 图像转置为 NCHW、把 HWIO 核转置为 OIHW：

from jax import lax out = lax.conv(jnp.transpose(img,[0,3,1,2]), # lhs = NCHW image tensor jnp.transpose(kernel,[3,2,0,1]), # rhs = OIHW conv kernel tensor (1, 1), # window strides 'SAME') # padding mode print("out shape: ", out.shape) print("First output channel:") plt.figure(figsize=(10,10)) plt.imshow(np.array(out)[0,0,:,:]);

lax.conv_with_general_padding则允许传入任意(low, high)padding 对，并支持 dilation：

out = lax.conv_with_general_padding( jnp.transpose(img,[0,3,1,2]), # lhs = NCHW image tensor jnp.transpose(kernel,[2,3,0,1]), # rhs = IOHW conv kernel tensor (1, 1), # window strides ((2,2),(2,2)), # general padding 2x2 (1,1), # lhs/image dilation (1,1)) # rhs/kernel dilation print("out shape: ", out.shape) print("First output channel:") plt.figure(figsize=(10,10)) plt.imshow(np.array(out)[0,0,:,:]);

Dimension Numbers：用字符串三元组声明维度布局

conv_general_dilated最关键的参数是维度说明三元组(Input Layout, Kernel Layout, Output Layout)，其中用到的字母含义如下：

• N：batch 维度
• H：空间高度
• W：空间宽度
• C：通道维度（输入/输出侧通用）
• I：核的输入通道维度
• O：核的输出通道维度

每个布局字符串长度必须等于张量秩，且 lhs/rhs/out 三者的空间字符集合必须一致。下面的例子为 NHWC 图像 + HWIO 核生成维度对象：

dn = lax.conv_dimension_numbers(img.shape, # only ndim matters, not shape kernel.shape, # only ndim matters, not shape ('NHWC', 'HWIO', 'NHWC')) # the important bit print(dn)

conv_dimension_numbers只关心张量的维数（ndim）而不关心具体形状（见 jax/_src/lax/convolution.py）：它会校验 lhs 与 rhs 维数相等、字符串长度正确、N/C/O/I各出现一次、且三个布局的空间字符集合一致，然后把字符串解析为ConvDimensionNumbers的轴排列对象，供底层 XLA 调用使用。

⚠️ 为展示 dimension numbers 的灵活性，下文lax.conv_general_dilated统一采用NHWC图像与HWIO核约定，与lax.conv的 NCHW/OIHW 默认约定不同。

场景一：SAME padding，无 stride，无 dilation

out = lax.conv_general_dilated(img, # lhs = image tensor kernel, # rhs = conv kernel tensor (1,1), # window strides 'SAME', # padding mode (1,1), # lhs/image dilation (1,1), # rhs/kernel dilation dn) # dimension_numbers = lhs, rhs, out dimension permutation print("out shape: ", out.shape) print("First output channel:") plt.figure(figsize=(10,10)) plt.imshow(np.array(out)[0,:,:,0]);

'SAME'会自动在空间维度两侧补齐，使输出空间尺寸与输入一致（本例输出仍为 200×198）。

场景二：VALID padding，无 stride，无 dilation

out = lax.conv_general_dilated(img, # lhs = image tensor kernel, # rhs = conv kernel tensor (1,1), # window strides 'VALID', # padding mode (1,1), # lhs/image dilation (1,1), # rhs/kernel dilation dn) # dimension_numbers = lhs, rhs, out dimension permutation print("out shape: ", out.shape, "DIFFERENT from above!") print("First output channel:") plt.figure(figsize=(10,10)) plt.imshow(np.array(out)[0,:,:,0]);

'VALID'不做任何填充，输出尺寸按空间尺寸 - 核尺寸 + 1收缩（本例 200-3+1=198、198-3+1=196），因此与'SAME'情况不同。

场景三：SAME padding，stride=(2,2)

out = lax.conv_general_dilated(img, # lhs = image tensor kernel, # rhs = conv kernel tensor (2,2), # window strides 'SAME', # padding mode (1,1), # lhs/image dilation (1,1), # rhs/kernel dilation dn) # dimension_numbers = lhs, rhs, out dimension permutation print("out shape: ", out.shape, " <-- half the size of above") plt.figure(figsize=(10,10)) print("First output channel:") plt.imshow(np.array(out)[0,:,:,0]);

窗口以 2 为步长滑动，配合'SAME'后输出空间尺寸约为输入的一半（100×99），这正是下采样（池化替代）的典型模式。

场景四：VALID padding + rhs 核 dilation —— 空洞卷积

out = lax.conv_general_dilated(img, # lhs = image tensor kernel, # rhs = conv kernel tensor (1,1), # window strides 'VALID', # padding mode (1,1), # lhs/image dilation (12,12), # rhs/kernel dilation dn) # dimension_numbers = lhs, rhs, out dimension permutation print("out shape: ", out.shape) plt.figure(figsize=(10,10)) print("First output channel:") plt.imshow(np.array(out)[0,:,:,0]);

rhs_dilation=(12,12)在核元素之间插入空洞（此处取值偏大，仅用于演示效果），等效于在不增加参数量的前提下扩大感受野，这就是空洞/膨胀卷积（Atrous Convolution）。

场景五：lhs 输入 dilation —— 转置卷积的基础

out = lax.conv_general_dilated(img, # lhs = image tensor kernel, # rhs = conv kernel tensor (1,1), # window strides ((0, 0), (0, 0)), # padding mode (2,2), # lhs/image dilation (1,1), # rhs/kernel dilation dn) # dimension_numbers = lhs, rhs, out dimension permutation print("out shape: ", out.shape, "<-- larger than original!") plt.figure(figsize=(10,10)) print("First output channel:") plt.imshow(np.array(out)[0,:,:,0]);

lhs_dilation=(2,2)在输入的空间轴元素之间插入零（即"上采样"输入），使输出尺寸大于原始输入——这正是转置卷积（Transposed Convolution / 分数步长卷积）的本质机制之一。

实战：用 conv_general_dilated 实现转置卷积

利用上面的 lhs dilation 技巧，可以手工实现转置卷积。其原理是：转置卷积 = 核旋转 180° + lhs 输入 dilation + 自定义输出 padding。下面的实现与 TensorFlow 的tf.nn.conv2d_transpose(img, kernel, (N,2*H,2*W,C), (1,2,2,1))等价：

# The following is equivalent to tensorflow: # N,H,W,C = img.shape # out = tf.nn.conv2d_transpose(img, kernel, (N,2*H,2*W,C), (1,2,2,1)) # transposed conv = 180deg kernel rotation plus LHS dilation # rotate kernel 180deg: kernel_rot = jnp.rot90(jnp.rot90(kernel, axes=(0,1)), axes=(0,1)) # need a custom output padding: padding = ((2, 1), (2, 1)) out = lax.conv_general_dilated(img, # lhs = image tensor kernel_rot, # rhs = conv kernel tensor (1,1), # window strides padding, # padding mode (2,2), # lhs/image dilation (1,1), # rhs/kernel dilation dn) # dimension_numbers = lhs, rhs, out dimension permutation print("out shape: ", out.shape, "<-- transposed_conv") plt.figure(figsize=(10,10)) print("First output channel:") plt.imshow(np.array(out)[0,:,:,0]);

注意这里 padding 从字符串换成了显式的((2, 1), (2, 1))——每个空间维度一个(low, high)对。原因见 jax/_src/lax/convolution.py：字符串 padding（'SAME'/'VALID'）在 lhs dilation 大于 1 时未实现，会直接抛出ValueError，必须显式指定每个维度的前后填充量，或者改用lax.conv_transpose。

如果不想手工旋转核、算 padding，JAX 还提供了现成的lax.conv_transpose便捷函数（源码见 jax/_src/lax/convolution.py）：它直接计算分数步长卷积，padding='SAME'/'VALID'会被解释为对应前向卷积的转置，也可传入显式 padding 对；transpose_kernel=True时自动翻转核的空间轴并交换输入/输出通道轴，使结果与 KerasConv2DTranspose等基于梯度的实现一致。它内置了NC/NHC/NHWC/NHWDC等各维数的默认布局。

1D 卷积：NWC 布局示例

conv_general_dilated并不局限于二维，一个简单的一维示例：

# 1D kernel - WIO layout kernel = jnp.array([[[1, 0, -1], [-1, 0, 1]], [[1, 1, 1], [-1, -1, -1]]], dtype=jnp.float32).transpose([2,1,0]) # 1D data - NWC layout data = np.zeros((1, 200, 2), dtype=jnp.float32) for i in range(2): for k in range(2): x = 35*i + 30 + 60*k data[0, x:x+30, k] = 1.0 print("in shapes:", data.shape, kernel.shape) plt.figure(figsize=(10,5)) plt.plot(data[0]); dn = lax.conv_dimension_numbers(data.shape, kernel.shape, ('NWC', 'WIO', 'NWC')) print(dn) out = lax.conv_general_dilated(data, # lhs = image tensor kernel, # rhs = conv kernel tensor (1,), # window strides 'SAME', # padding mode (1,), # lhs/image dilation (1,), # rhs/kernel dilation dn) # dimension_numbers = lhs, rhs, out dimension permutation print("out shape: ", out.shape) plt.figure(figsize=(10,5)) plt.plot(out[0]);

这里数据使用NWC布局（batch、宽度、通道），核使用WIO布局（宽度、输入通道、输出通道），两个通道分别送入不同的脉冲序列，卷积后得到包含两个输出通道的结果。注意 1D 场景下所有 stride/dilation 元组长度相应变为 1。

3D 卷积：NHWDC 布局示例

同理，三维卷积只需把布局推广到NHWDC/HWDIO：

import matplotlib as mpl # Random 3D kernel - HWDIO layout kernel = jnp.array([ [[0, 0, 0], [0, 1, 0], [0, 0, 0]], [[0, -1, 0], [-1, 0, -1], [0, -1, 0]], [[0, 0, 0], [0, 1, 0], [0, 0, 0]]], dtype=jnp.float32)[:, :, :, jnp.newaxis, jnp.newaxis] # 3D data - NHWDC layout data = jnp.zeros((1, 30, 30, 30, 1), dtype=jnp.float32) x, y, z = np.mgrid[0:1:30j, 0:1:30j, 0:1:30j] data += (jnp.sin(2*x*jnp.pi)*jnp.cos(2*y*jnp.pi)*jnp.cos(2*z*jnp.pi))[None,:,:,:,None] print("in shapes:", data.shape, kernel.shape) dn = lax.conv_dimension_numbers(data.shape, kernel.shape, ('NHWDC', 'HWDIO', 'NHWDC')) print(dn) out = lax.conv_general_dilated(data, # lhs = image tensor kernel, # rhs = conv kernel tensor (1,1,1), # window strides 'SAME', # padding mode (1,1,1), # lhs/image dilation (1,1,1), # rhs/kernel dilation dn) # dimension_numbers print("out shape: ", out.shape) # Make some simple 3d density plots: from mpl_toolkits.mplot3d import Axes3D def make_alpha(cmap): my_cmap = cmap(jnp.arange(cmap.N)) my_cmap[:,-1] = jnp.linspace(0, 1, cmap.N)**3 return mpl.colors.ListedColormap(my_cmap) my_cmap = make_alpha(plt.cm.viridis) fig = plt.figure() ax = fig.add_subplot(projection='3d') ax.scatter(x.ravel(), y.ravel(), z.ravel(), c=data.ravel(), cmap=my_cmap) ax.axis('off') ax.set_title('input') fig = plt.figure() ax = fig.add_subplot(projection='3d') ax.scatter(x.ravel(), y.ravel(), z.ravel(), c=out.ravel(), cmap=my_cmap) ax.axis('off') ax.set_title('3D conv output');

三维核为 3×3×3（HWDIO，输入输出通道各 1），数据为 30×30×30 的三维体数据（NHWDC），mode='SAME'下输出保持 30×30×30 的空间尺寸。

深入源码：conv_general_dilated 的完整参数面

conv_general_dilated的完整签名（见 jax/_src/lax/convolution.py）比示例中用到的参数更丰富，完整参数与语义如下：

参数	含义	默认值
lhs	秩为n+2的输入张量	必填
rhs	秩为n+2的卷积核张量	必填
window_strides	长度为n的窗口步长序列	必填
padding	'SAME'/'SAME_LOWER'/'VALID'字符串，或长度为n的(low, high)对序列	必填
lhs_dilation	输入各空间维的 dilation（即转置卷积机制）	(1,)*n
rhs_dilation	核各空间维的 dilation（即空洞卷积机制）	(1,)*n
dimension_numbers	None、ConvDimensionNumbers或(lhs_spec, rhs_spec, out_spec)字符串三元组	None→('NCHW','OIHW','NCHW')（2D）
feature_group_count	特征分组数，用于深度可分离/分组卷积	1
batch_group_count	batch 分组数	1
precision	None、Precision.DEFAULT/HIGH/HIGHEST或字符串（如'highest'/'fastest'），也可传二元组分别指定 lhs/rhs 精度	None
preferred_element_type	累加与返回的目标数据类型	None（按输入默认）

几个值得注意的源码级细节：

• 默认布局：当dimension_numbers=None时，二维卷积默认采用('NCHW', 'OIHW', 'NCHW')（见 jax/_src/lax/convolution.py），与 TensorFlowConv2D的('NHWC','HWIO','NHWC')不同，跨框架迁移时需显式声明。
• padding 字符串语义：'SAME'与'SAME_LOWER'都把输出补到与输入等大，差别只在奇数 padding 时多余的一行/列加在末尾（'SAME'）还是开头（'SAME_LOWER'）。传入字符串 padding 时，源码会结合 lhs 形状、有效核尺寸（考虑rhs_dilation）与 stride 计算出具体 pad 值（jax/_src/lax/convolution.py）。
• stride 与空间字符的绑定：字符串形式的 dimension numbers 中，window_strides[i]按rhs_spec中除'I'/'O'外第一个空间字符出现的顺序与维度匹配（jax/_src/lax/convolution.py），因此使用非标准布局时务必核对 stride 顺序。
• 校验严格：conv_dimension_numbers会拒绝 lhs/rhs 维数不等、字符串长度不符、N/C/O/I出现次数不为 1、存在重复字符或三个布局空间字符集合不一致的输入（jax/_src/lax/convolution.py），这也是"张量布局声明错误"最常见的报错来源。
• 与自动微分的配合：conv_general_dilated是标准 lax 原语，可被jax.grad、jax.jit、jax.vmap直接组合使用；源码中_conv_general_vjp_lhs_padding与_conv_general_vjp_rhs_padding（jax/_src/lax/convolution.py）分别实现了对 lhs 与 rhs 的梯度 padding 推导，这正是"卷积层反向传播自动可用"的底层保证。

小结：如何选择正确的卷积接口

• 一维信号处理（平滑、滤波、相关分析）→jax.numpy.convolve/jax.numpy.correlate；
• N 维/二维图像处理，追求 SciPy 兼容语义 →jax.scipy.signal.convolve/convolve2d，大核大输入可试method='fft'；
• 深度学习中的批量卷积，需要自由控制 stride、padding、dilation 与维度布局 →lax.conv_general_dilated（配合lax.conv_dimension_numbers声明布局），需要转置卷积时优先用lax.conv_transpose。

无论选择哪一层，最终的计算都会落到 jax/_src/lax/convolution.py 中的conv_general_dilated原语上，而它直接封装 XLA 的Conv算子，因而天然支持 CPU、GPU 与 TPU 上的统一执行与自动微分。

【免费下载链接】jaxComposable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/gh_mirrors/jax/jax