★ FEATURED ARTICLE 最近一年我把不少模型从PyTorch搬到了JAX起因也是看了几个项目和论文的官方实现——训练代码简洁到让人怀疑人生自动微分和编译加速是天然集成的。真正动手写了一段时间之后我的感受是JAX的确不是一个换个框架写同样代码的替代品它的函数式范式会倒逼你把计算流程想得比之前更清楚。这篇文章就用一个非常经典的模型——图卷积网络GCN——来串一遍JAX的核心用法。内容会从JAX的基本功开始一直写到两层GCN的完整实现和训练包含能直接复制运行的代码以及我在踩坑之后才想明白的几个关键点。这篇内容适合谁如果你是刚接触JAX、之前主要用PyTorch或者纯NumPy想搞清楚jit、grad、vmap到底是什么以及它们如何组合在一起工作或者你想在一个真实的小型图数据集上跑通GCN但不想被TensorFlow那套生态绑住这篇都很合适。读完之后你至少能自己动手复现一个GCN并且对JAX的编程习惯有一个整体把握。1. JAX的底层气质先想清楚它为什么长这样1.1 它是NumPy但比NumPy更纯净JAX最表面的形态是jax.numpy接口长得和NumPy几乎一样。第一次跑jnp.array([1,2,3])的时候你会觉得毫无门槛随便写两行矩阵运算手感完全一致。但核心区别藏在数组不可变这个规定里。在NumPy或者PyTorch里你习惯了x[0] 1原地修改或者params.grad自动累积梯度这些在JAX里都不存在——数组一旦创建就不能改想更新某个位置只能通过.at操作符返回一个新数组。import jax.numpy as jnp x jnp.array([1.0, 2.0, 3.0]) # x[0] 10.0 # 报错TypeError y x.at[0].set(10.0) print(x) # [1. 2. 3.]x没变 print(y) # [10. 2. 3.]y是新数组这个设计初看很啰嗦但它是整个JAX能实现jit和grad的前提。因为程序里没有共享可变状态编译器才能放心大胆地做表达式重写、算子融合、甚至是并行切分。换句话说JAX选择了一条信任程序纯度的路计算纯粹是由输入到输出的函数映射所有副作用都被剔除。你写的每一段训练逻辑最后都会被编译成一张无状态的计算图。我刚从PyTorch转过来的时候最难受的就是这件事。PyTorch里模型参数就是一个Tensor你随时可以param.grad累加、optimizer.step()原地更新到了JAX里参数就是普通的数据结构用jax.tree_util管理更新是老参数 负学习率×梯度这样一个生成新结构的过程。你会被迫写一个显式的update函数。这件事一开始觉得烦但后来发现有个实际好处训练循环是确定性的不会因为某个状态被意外修改产生隐性问题复现实验变得异常简单。1.2grad不是装饰器而是函数变换器PyTorch里做反向传播要用loss.backward()然后参数上的.grad被填好再交给优化器。JAX里没有这一套隐式状态它把求梯度变成一个显式的函数变换grad(f)接受一个函数返回这个函数的梯度版本。import jax import jax.numpy as jnp def f(x): return jnp.sum(x ** 2) g jax.grad(f) print(g(jnp.array([1.0, 2.0, 3.0]))) # [2. 4. 6.]注意几点grad默认只对第一个参数求导函数必须返回标量否则你不知道梯度应该顺着哪个方向传。如果你有一个函数要同时对多个参数求导用argnums指定位置即可。如果某个参数是嵌套结构——比如一个字典的字典——grad也能正常返回相同结构的梯度。这一点非常重要因为GCN的两个权重矩阵和两个偏置我通常就装在一个简单的列表里grad照样能算出结构完全一致的梯度。另外还有jax.value_and_grad它在求梯度的同时把函数值也返回回来。训练里我们通常既想看loss又需要梯度用这个变换就可以只写一遍前向计算省得forward被调用两次。loss_and_grad jax.value_and_grad(loss_fn) loss_val, grads loss_and_grad(params, x, A_hat, labels, mask)这里loss_fn的签名是loss_fn(params, ...)我习惯让params永远放在第一位这样value_and_grad的默认行为正好指向它不用额外设置argnums。1.3jit到底加速了什么jit是JAX的即时编译器底层通过XLA把Python函数翻译成高效的、可以融合算子的计算图。用起来就是在函数定义上加一个装饰器或者对已经定义的函数调用一次jax.jit(func)。jax.jit def forward(params, x, A_hat): h A_hat x params[0] params[1] h jax.nn.relu(h) h A_hat h params[2] params[3] return h第一次调用这个函数时XLA会把Python层的计算流程编译成底层指令这个编译过程通常需要几百毫秒甚至更久之后相同输入形状的调用都走编译好的代码速度会有明显提升。我的经验是如果模型不算太复杂jit带来的收益在小数据集上不一定直观但一旦batch size或者图的规模上来了算子融合和内存分配的优化就非常明显跑大图的时候能省一半以上的时间。jit也有自己的脾气最大的限制是它要求函数的控制流和数据形状最好在编译期就能确定。如果函数内部写if x.sum() 0这种依赖数组值的Python分支jit会出错或者表现奇怪。遇到这种情况需要用jax.lax.cond这类带条件语义的原语替代。GCN的前向计算本质是矩阵乘法链几乎碰不到这种问题这也是我挑它来做示范的原因之一。1.4 随机数大概是新人最不习惯的地方JAX的随机数设计是另一个劝退点。PyTorch里torch.randn直接按全局随机状态生成JAX则要求你显式携带一个随机密钥key每次生成随机数之前必须用jax.random.split来分裂出新key。否则你会发现两次调用生成的是完全一样的数组——这不是bug是设计。key jax.random.PRNGKey(42) # 需要两个正态数组时 key1, key2 jax.random.split(key, 2) w1 jax.random.normal(key1, (input_dim, hidden_dim)) w2 jax.random.normal(key2, (hidden_dim, output_dim))这种设计一开始很烦尤其是写训练循环的时候模型里一旦有dropout就必须在每轮迭代里重新split一个key然后把它传进forward。但好处是随机性可以被完全复现只要保存好初始key整个实验序列都能重演。GCN的初始化代码我会在第三节完整给出到时候你会看到所有随机数组都是这样从key里长出来的逻辑非常透明。2. GCN的前向公式与邻接矩阵归一化原理2.1 图卷积到底卷积了什么GCN是图神经网络里最经典、最朴素的一支。它的设计动机非常简单普通卷积网络处理的是规则网格图像、语音而图上的节点没有固定顺序也没有平移不变性所以我们没法定义卷积核滑过邻域这件事。Kipf那篇经典工作给出的替代方案是每一层先通过邻接矩阵聚合邻居的特征再乘上一个可学习的权重矩阵。它的核心公式是H^(l1) σ( Â · H^(l) · W^(l))其中Â是经过归一化的邻接矩阵H是节点特征W是权重。这个公式看起来就是一个矩阵连乘没有任何复杂的结构所以用JAX来实现非常自然——jnp.dot和运算符就能覆盖全部计算。动态上理解这个公式邻接矩阵A里A_ij 1表示节点i和节点j之间有边那么A·H就是每个节点把自己邻居的特征拿过来相加。这个操作的意义是每个节点的表示不再只看自己而是看自己和周围邻居的组合。一个节点如果邻居很多它的聚合结果就很大数值不稳定如果邻居很少又显得更新量很小。为了公平起见需要用度矩阵D来归一化。2.2 为什么是 D^(-1/2)·A·D^(-1/2) 而不是 D^(-1)·A归一化有两种常见做法。第一种是行归一化D^(-1)A效果是每行之和为1所有邻居的特征做加权平均。第二种是Kipf用的对称归一化D^(-1/2)AD^(-1/2)效果是每个节点更新时邻居的贡献会被两端的度同时缩放。后者有非常明确的谱域意义它对应一个归一化图拉普拉斯算子上的线性滤波器数学上能推导出它近似一种低通滤波让相似的节点表示更相似。从实现角度讲对称归一化的计算方式也很优雅def normalize_adjacency(A): A_tilde A jnp.eye(A.shape[0]) deg jnp.sum(A_tilde, axis1) deg_inv_sqrt jnp.power(deg, -0.5) deg_inv_sqrt jnp.where(jnp.isinf(deg_inv_sqrt), 0.0, deg_inv_sqrt) D_inv_sqrt jnp.diag(deg_inv_sqrt) return D_inv_sqrt A_tilde D_inv_sqrt这里加单位阵的作用是给每个节点加自环。如果不加自环节点在聚合时看不到自己的特征只能看到邻居这样关于自身的信息在每一层都会被稀释深层模型很容易丢失节点身份信息。加了自环相当于默认我自己也是我邻居之一这被后来的无数实验证实是GCN的标配。jnp.where(jnp.isinf(...))是为了防止孤立节点度为0在取-0.5次方时出现无穷大。你可能注意到jnp.power(deg, -0.5)中的deg是零的时候结果会是无穷大在转换成对角矩阵后会生成NaN。所以先把它替换成0既保证了数值安全也确保孤立节点不会贡献任何聚合信息。这一步在PyTorch里你得费点心思但在JAX里jnp.where配合jnp.isinf就是一次优雅的向量化操作。2.3 GCN层数的2层就够直觉实现GCN之前你还应该知道一个经验结论大多数标准场景下2层GCN已经够了。原因主要是过平滑问题——图卷积本质上是一种平滑操作每一层都在让邻居节点的表示趋于一致层数越多所有节点的表示越发趋同反而损害分类精度。所以你在各种论文里看到的基线方案基本就是2层输入 - 隐藏层ReLU - 输出层。用我这个项目做例子输入维度是特征维度隐藏层设成16输出维度是类别数参数总量大概只有2×输入维×16 16×类别数小到CPU上跑几百轮都毫无压力。这个模型小但效果好的特质让它特别适合作为JAX的练手项目——你不需要调用复杂的自定义算子也不需要处理分布式训练前向和反向都干净利落但又能把grad、jit、随机数、参数更新全部串起来。3. 用JAX实现一个两层GCN的完整过程3.1 准备数据图结构、特征和标签我这里用一个经典的论文引用网络数据集来演示它的规模很小适合在CPU上快速验证。数据集包含约2700个节点每个节点是一篇论文特征是对应词袋向量0/1表示是否包含某个关键词标签是论文的类别节点之间的边表示引用关系。目标是仅用一部分节点的标签来训练模型然后预测剩下的节点。由于JAX本身不带数据加载器预处理部分可以直接用常规Python处理。通常这种数据集会提供三个数组features节点特征矩阵、adjacency邻接矩阵、labels标签向量还有train_mask/val_mask/test_mask分别标记哪些节点用于训练、验证、测试。import numpy as np import jax.numpy as jnp # 假设已经从数据文件加载到了numpy数组 # features: (N, F) # adjacency: (N, N) 稀疏或稠密 # labels: (N,) 整数 # train_mask: (N,) bool N features.shape[0] F features.shape[1] C int(labels.max()) 1 A np.asarray(adjacency.todense(), dtypenp.float32) # 小图可以直接转稠密 X jnp.array(features, dtypejnp.float32) Y jnp.array(labels, dtypejnp.int32) A_hat jnp.array(normalize_adjacency(jnp.array(A, dtypejnp.float32)), dtypejnp.float32) train_mask jnp.array(train_mask, dtypejbool) val_mask jnp.array(val_mask, dtypejbool) test_mask jnp.array(test_mask, dtypejbool)如果你用的是稀疏邻接矩阵JAX对稀疏的支持目前在CPU上还不够顺手所以小图我建议直接稠密化。2700×2700的矩阵才几十MB完全没问题。对大图希望用稀疏实现的话我后面会简单提一下方向但本篇文章以理解原理为主。3.2 初始化参数JAX风格的参数管理参数初始化在JAX里有一个需要强迫自己记住的原则参数必须是纯数据结构不能挂任何隐藏状态。我把两个层的权重和偏置放在一个简单的列表里def init_params(key, input_dim, hidden_dim, output_dim): key1, key2 jax.random.split(key, 2) w1 jax.random.normal(key1, (input_dim, hidden_dim)) * jnp.sqrt(2.0 / input_dim) b1 jnp.zeros((hidden_dim,)) w2 jax.random.normal(key2, (hidden_dim, output_dim)) * jnp.sqrt(2.0 / hidden_dim) b2 jnp.zeros((output_dim,)) return [w1, b1, w2, b2]缩放因子sqrt(2/input_dim)是参考了ReLU网络的初始化经验目的是让前向传播的方差保持稳定。如果你不缩放网络在深层时数值很容易爆炸GCN虽然只有两层但最好还是从一开始就养成好习惯。init_params返回的是一个Python列表。JAX的jit、grad、value_and_grad可以自动递归处理列表、字典、元组这些结构这也是jax.tree_util在背后做的事。你可以放心地把这个列表传给梯度函数返回的梯度依然是一个同构的列表。3.3 前向传播两个矩阵乘法就够前向传播函数定义为def forward(params, X, A_hat): w1, b1, w2, b2 params h A_hat X w1 b1 h jax.nn.relu(h) logits A_hat h w2 b2 return logits这里的三次矩阵乘法分别对应输入特征在图上的第一次平滑传播聚合邻居信息、隐藏层到输出层的第二次平滑传播。你可以观察到这个实现没有显式的邻居采样注意力机制边特征融合它就是最朴素的聚合加变换。恰恰因为朴素它才能在JAX里被jit完美加速——整个计算就是一个纯线性代数表达式XLA可以无压力地做算子融合。有一点值得注意这里A_hat X w1先做A_hat X再做结果与w1的乘法。从计算顺序上讲先和图矩阵乘会让聚合发生在特征变换之前这与一些实现先X w1再做聚合在数学上其实等价因为矩阵乘法满足结合律。但两种顺序的浮点误差和效率可能有细微差异我在实践中习惯让A_hat在最左边符合先聚合再变换的直觉。3.4 损失函数与训练循环GCN训练一般用交叉熵。因为只有带标签的节点才有监督信号所以要把mask应用进去。def loss_fn(params, X, A_hat, Y, mask): logits forward(params, X, A_hat) log_probs jax.nn.log_softmax(logits) label_onehot jax.nn.one_hot(Y, num_classeslogits.shape[-1]) loss -jnp.mean(jnp.sum(log_probs * label_onehot, axis-1) * mask) return loss这里的mask是布尔向量True表示该节点参与loss计算。乘法的效果是让未标记节点贡献为0jnp.mean分母上会包含0贡献的元素所以严格来说这不是标准的平均CE不过在图半监督场景下大家一般都用这个写法因为它简单且避免了一些除零判断。如果你想做严格的平均可以改成loss.sum() / mask.sum()差别很小。接下来是训练循环的主角——update函数。它做了四件事计算loss、计算梯度、更新参数、返回新的loss值。整个过程完全无副作用每次调用返回新的参数结构。jax.jit def update(params, X, A_hat, Y, mask, lr): loss_val, grads jax.value_and_grad(loss_fn)(params, X, A_hat, Y, mask) new_params [ p - lr * g for p, g in zip(params, grads) ] return new_params, loss_val你可能会问参数更新里用了Python的列表推导式jit能处理吗能。因为params和grads都是长度固定的列表推导式在编译阶段就被展开了XLA清楚地知道每一步操作是什么。这里lr是一个Python浮点数它会作为编译期的常量参数被烘焙进图里。所以如果你每轮都改变学习率会导致JIT重新编译。实际训练中我们的学习率通常是固定值或按epoch调整不是每步都变这个问题不大。完整训练循环写出来大概是key jax.random.PRNGKey(0) params init_params(key, F, 16, C) lr 0.01 for epoch in range(200): params, loss_val update(params, X, A_hat, Y, train_mask, lr) if epoch % 20 0: acc evaluate(params, X, A_hat, Y, test_mask) print(fepoch {epoch}: loss{loss_val:.4f} test_acc{acc:.4f})evaluate函数可以顺便演示一下vmap或者argmaxdef evaluate(params, X, A_hat, Y, mask): logits forward(params, X, A_hat) preds jnp.argmax(logits, axis-1) correct jnp.sum((preds Y) * mask) total jnp.sum(mask) return correct / total这一套全跑完之后在标准引文网络数据集的实验上你大概率能看到测试准确率在0.75到0.80之间。不同随机种子和初始化下会有几个点的波动但整体趋势是稳定的。如果你发现结果差得离谱第一件事检查有没有把自环加上第二件事检查学习率是否太大——我见过有人设成0.5直接发散到NaN。4. 实际训练过程中踩过的坑与排查链路4.1 JIT的隐式重新编译让训练越跑越慢我第一次用JAX跑图模型时遇到过一个很隐蔽的性能问题训练前20轮正常后面突然每一步都变慢而且CPU风扇狂转。排查了半天发现问题出在A_hat的形状上——训练集合大小不变按理说形状不会变。但实际上我的数据预处理把A_hat做成了稠密矩阵而当时代码里用了np.asarray(adjacency.todense())todense()在某些稀疏后端上会返回一个奇怪的矩阵形状元信息导致传给jit函数时每次都被当作新的输入签名于是每个step都在重新编译。排查思路建议分三步第一步把update函数里的jax.jit临时去掉看速度是否正常。如果去掉后恢复正常基本确定是编译开销。第二步在update函数外打印A_hat.shape和A_hat.dtype确认这两者是稳定的。JIT的缓存键基于输入数组的形状和dtype任何一个变化都会触发重编译。第三步如果确实有某个输入在反复变化形状把它固定成统一形状或者把它移出jit函数只对内部矩阵乘法做预计算。这个坑最麻烦的地方在于JAX不会给你任何报错它只是安静地重编译而你觉得程序没坏只是慢得诡异。希望这里写出来能帮你省掉一个晚上的排查时间。4.2random.key管理失当导致Dropout失效我在另一个图模型里用了dropout当时天真地以为像PyTorch一样每个step自动有不同的随机行为。结果训练准确率怎么都上不去后来才发现我在每次迭代里把同一个key传进了forward而jax.random.bernoulli(key, ...)对同一个key永远返回同一个mask所以dropout实际上变成了一个固定比例的确定性丢弃而非随机丢弃。正确的做法是用jax.random.split产生新的key并把新key的一部分用于dropout另一部分保留给下一次迭代jax.jit def update_with_dropout(params, X, A_hat, Y, mask, key, lr): key, subkey jax.random.split(key) ... logits forward(params, X, A_hat, subkey) ...也就是说key是训练循环里必须显式流转的状态就像参数一样。GCN本身没有dropout时不会遇到这个问题但你一旦开始尝试正则化这个坑会准时出现。4.3 过平滑现象深GCN反而更差我在验证模型时顺手把层数从2层加到4层、6层结果测试准确率不升反降。这不是JAX的问题而是图卷积的天然特性每多一层每个节点就更倾向于吸收更远邻居的信息节点表示彼此趋同最终变成所有节点一个模子。后来我专门看了节点表示的相似度确实是一层比一层高这印证了过平滑的判断。所以在使用GCN时一个实用的建议是不要盲目堆深度。如果你的任务需要更多层的感受野尝试用残差连接、JK-Net、或者GAT这类更复杂的架构不要指望单纯加深GCN能带来收益。用JAX实现残差连接也很简单就是把上一层的表示和当前层的输出加在一起再激活一行代码的事。4.4 特征矩阵的数值范围某些数据集的特征是词袋0/1向量数值范围很小训练很稳定。但如果你换到带TF-IDF权重的特征范围差异会很大前向传播里可能出现极大值经过softmax后损失变成NaN。遇到这种情况我的经验是检查A_hat是否真的归一化正确以及输入特征做一次行归一化或者log变换。JAX里做这个非常自然X X / jnp.linalg.norm(X, axis1, keepdimsTrue)4.5 为什么我在代码里避免使用Python原生循环做邻居聚合如果你之前在PyTorch里写过GNN可能会习惯用for循环对每个节点手动找邻居。在JAX里这种方式在jit下很难高效编译因为循环次数取决于节点度是运行时才确定的值XLA通常无法把动态循环展开成高效指令。正确的姿势是把整个图的操作向量化成矩阵乘法——用A_hat H一次性计算所有节点的聚合结果。这正是GCN公式本身已经做好的事。如果将来你需要实现GraphSAGE那种邻居采样算法推荐的方向是用jax.lax.scan替代Python的for循环或者在jit外部用Python循环逐批处理。用vmap也可以对批量节点应用同一个聚合函数这是JAX进行图模型优化的通用工具箱。5. 进一步优化JIT加速可视化、vmap、以及更真实的训练体验5.1 用block_until_ready看到真实的JIT收益JAX是异步执行的你写loss_val update(...)的时候计算可能还没真正跑完只是被放进了队列。如果你用Python的time模块测速测到的可能是提交时间而不是实际计算时间这会让性能测试产生严重误判。解决办法是在计时前调用jax.block_until_ready。import time start time.time() for _ in range(100): params, loss_val update(params, X, A_hat, Y, train_mask, lr) loss_val.block_until_ready() elapsed time.time() - start print(f100 steps: {elapsed:.2f}s)第一次调用update时会包含JIT编译时间所以你在测速时最好先跑一次预热再正式开始计时。我在比较jit开和关的性能差距时踩过这个坑第一次跑出来的时间大得离谱后来才知道是编译开销。5.2vmap在GCN场景中的实际价值虽然GCN的前向本身是全局矩阵乘法不需要vmap参与但在模型调优和诊断阶段vmap非常有用。比如你想对同一组参数快速测试10个不同随机种子的初始化效果可以把初始化和前向都映射到一个向量化函数上from functools import partial def init_and_forward(key, X, A_hat): params init_params(key, F, 16, C) return forward(params, X, A_hat) keys jax.random.split(jax.random.PRNGKey(0), 10) logits_batch jax.vmap(partial(init_and_forward, XX, A_hatA_hat))(keys)vmap自动把keys的第一个维度当成batch维其他参数保持不变。它能让你在写批量实验时避免手写Python循环而且这种向量化方式在jit下通常能获得更好的编译结果。这个技巧在你做超参数扫参的时候非常实用建议收藏。5.3 训练过程的动态日志与早停我还想分享一个很实在的经验哪怕模型很小也建议写一个基于验证集准确率的早停机制。GCN的loss曲线在标准引文数据集上经常是先快速下降然后在一个平台期轻微波动如果没有早停最后取到的checkpoint未必是验证集上最好的。JAX没有现成的EarlyStopping回调但实现起来只需要维护一个最优参数的副本和计数变量best_params params best_val_acc 0.0 patience 20 wait 0 for epoch in range(500): params, loss_val update(params, X, A_hat, Y, train_mask, lr) val_acc evaluate(params, X, A_hat, Y, val_mask) if val_acc best_val_acc: best_val_acc val_acc best_params params wait 0 else: wait 1 if wait patience: break这个模式在几乎所有JAX项目里都能复用。它的价值在于避免你拿测试集反复调参造成信息泄露同时在模型接近过平滑时及时停下来而不是等到所有节点都变成同一个向量。5.4 模型评估的一个细节mask的数据类型最后提醒一个容易忽略的细节mask数组的dtype必须是布尔型。如果你从数据集里读到的mask是整数0/1jnp.sum(mask)也能算但一旦参与乘法整型mask会把浮点梯度变成全零或者意外截断浮点数。我在一次实验里就因为这个把测试准确率算成了0排查了半天。稳妥的做法是读取数据后统一做一次类型转换mask jnp.array(mask, dtypejbool)6. 一些额外的个人体会从PyTorch迁到JAX最值回票价的不是某个具体算子的快慢而是一种思维上的瘦身。你被迫把训练过程拆成init、forward、loss、update这么几个纯函数任何中间状态都无处躲藏实验逻辑因此变得极度透明。GCN作为一个只需十行前向代码的模型恰恰能让你把注意力完全放在JAX本身的特性上而不是被复杂的网络结构干扰。如果你已经顺利跑通了上面的代码我建议你接着做三件事一是把hidden_dim从16调整到64、128观察训练速度和精度的变化二是把dropout加进两层之间用jax.random.split正确管理key看看正则化对最终准确率的影响三是把数据集换成一个更大规模的引文网络感受一下稠密A_hat带来的内存压力这时你可以开始研究JAX对稀疏矩阵的支持或者用vmap配合邻居采样来实现更灵活的训练方式。这三步做完你对JAX和GCN的掌握都会再上一个台阶。
阅读完成 · 觉得有帮助？
