3、自动微分与梯度计算：grad函数、jacfwd与jacrev、vmap向量化映射

好，咱们今天聊点硬核的。自动微分，说白了就是让框架帮我们把导数算明白。你想想看，在自动驾驶感知里，不管是目标检测还是语义分割，最后都要靠梯度下降来优化模型参数。如果梯度算错了，那整个模型就废了。

JAX 的自动微分做得非常优雅。我个人习惯把它分成三个层次来理解：标量梯度、雅可比矩阵、向量化加速。咱们一个一个来。

3.1 grad：最基础的梯度计算

grad 函数是 JAX 自动微分的入口。你给它一个函数，它返回这个函数的梯度函数。就这么简单。

举个例子，假设我们要计算一个简单的感知机层的前向传播和梯度：

import jax
import jax.numpy as jnp

def linear_layer(params, x):
    w, b = params
    return jnp.dot(w, x) + b

def loss_fn(params, x, y_true):
    y_pred = linear_layer(params, x)
    return jnp.mean((y_pred - y_true) ** 2)

# 初始化参数
key = jax.random.PRNGKey(42)
w = jax.random.normal(key, (3, 4))
b = jnp.zeros(3)
params = (w, b)

# 生成数据
x = jnp.ones(4)
y_true = jnp.array([1.0, 0.0, -1.0])

# 计算梯度
grad_loss = jax.grad(loss_fn)
grads = grad_loss(params, x, y_true)
print(grads)

嗯，这里要注意：grad 默认只对第一个参数求导。如果你的损失函数有多个参数，你需要用 argnums 指定对哪个参数求导。我在项目中遇到过这种情况——一个多任务学习的模型，需要对不同任务的头部分别求导，这时候 argnums 就派上用场了。

小技巧： 如果你需要同时计算函数值和梯度，可以用 jax.value_and_grad。这样一次前向传播就能拿到两个结果，省掉一次计算。在自动驾驶的实时推理场景里，这点效率提升很关键。

3.2 jacfwd 与 jacrev：雅可比矩阵的计算

说实话，grad 只能处理标量输出。但很多时候，我们需要的是输出向量对输入向量的导数——也就是雅可比矩阵。比如在车辆轨迹预测中，我们需要知道预测轨迹点对输入特征的敏感度。

JAX 提供了两种计算雅可比矩阵的方式：jacfwd（前向模式）和 jacrev（反向模式）。

为什么要有两种？这涉及到计算复杂度的问题。我简单解释一下：

  • jacfwd：适合「输入维度小，输出维度大」的场景。比如输入是 10 维，输出是 1000 维。
  • jacrev：适合「输入维度大，输出维度小」的场景。比如输入是 1000 维，输出是 10 维。

来看个实际例子。假设我们在做激光雷达点云的体素特征提取：

def voxel_feature_extractor(points, weights):
    # points: (N, 3)  N个点，每个点3维坐标
    # weights: (3, 64) 特征提取权重
    features = jnp.dot(points, weights)  # (N, 64)
    return jnp.mean(features, axis=0)    # (64,)

# 假设有1000个点
points = jax.random.normal(key, (1000, 3))
weights = jax.random.normal(key, (3, 64))

# 计算雅可比矩阵
jacobian_fwd = jax.jacfwd(voxel_feature_extractor, argnums=0)(points, weights)
jacobian_rev = jax.jacrev(voxel_feature_extractor, argnums=0)(points, weights)

# 验证结果是否一致
print(jnp.allclose(jacobian_fwd, jacobian_rev))  # True
我曾经踩过的坑： 有一次我在做多传感器融合的梯度计算，输入维度是 2000+，输出只有 6 维（位姿）。我一开始用了 jacfwd，结果显存直接爆了。后来换成 jacrev，内存占用降了 10 倍。所以选对模式真的很重要。

3.3 vmap：向量化映射的艺术

好了，现在你有了梯度，有了雅可比矩阵。但还有一个问题：效率。

在自动驾驶中，我们经常需要对一个 batch 的数据做同样的操作。比如同时处理 64 张图像，或者同时计算 100 个候选轨迹的梯度。最笨的办法是用 for 循环，但那样太慢了。

vmap 就是来解决这个问题的。它能把一个处理单个样本的函数，自动变成处理 batch 的函数。而且这个过程是自动向量化的，不需要你手动写 batch 维度。

来看个对比：

def single_sample_loss(params, image):
    # 处理单张图像
    features = jnp.dot(image, params['w']) + params['b']
    return jnp.mean(features ** 2)

# 生成 batch 数据
batch_size = 64
images = jax.random.normal(key, (batch_size, 224, 224, 3))
params = {'w': jax.random.normal(key, (3, 64)), 'b': jnp.zeros(64)}

# 方法1：手动 for 循环（慢）
losses_loop = jnp.array([single_sample_loss(params, img) for img in images])

# 方法2：使用 vmap（快）
batched_loss_fn = jax.vmap(single_sample_loss, in_axes=(None, 0))
losses_vmap = batched_loss_fn(params, images)

print(jnp.allclose(losses_loop, losses_vmap))  # True

这里 in_axes=(None, 0) 的意思是：第一个参数（params）不 batch 化，第二个参数（images）在第 0 维 batch 化。你想想看，这个机制在自动驾驶的多任务学习中特别有用——共享的 backbone 参数不 batch，不同任务的 head 各自 batch。

3.4 三者的组合拳

在实际项目中，grad、jacfwd/jacrev、vmap 往往是组合使用的。我画了一张图来展示它们的关系：

JAX 自动微分三件套 函数 f(x) grad 标量梯度 ∇f jacfwd/jacrev 雅可比矩阵 J vmap 向量化映射 实际组合使用场景 vmap(grad(loss_fn)) → 批量计算梯度 vmap(jacfwd(fn)) → 批量计算雅可比矩阵 选择依据：输入/输出维度

你看，这三者的关系其实很清晰：grad 是基础，jacfwd/jacrev 是扩展，vmap 是加速器。在实际的自动驾驶感知代码中，我经常这样组合使用：

# 完整的训练步骤
def compute_batch_grads(params, batch_images, batch_labels):
    # 定义单样本损失
    def single_loss(params, img, label):
        pred = model(params, img)
        return jnp.mean((pred - label) ** 2)
    
    # 批量计算梯度
    batch_grad_fn = jax.vmap(
        jax.grad(single_loss), 
        in_axes=(None, 0, 0)
    )
    
    return batch_grad_fn(params, batch_images, batch_labels)

# 使用示例
batch_grads = compute_batch_grads(params, images, labels)
核心要点：
  • grad 用于标量输出函数的梯度计算
  • jacfwd 适合「小输入、大输出」场景
  • jacrev 适合「大输入、小输出」场景
  • vmap 自动向量化，替代手动 batch 循环
  • 三者可以任意组合，实现高效的批量梯度计算

说实话，刚开始接触 JAX 的时候，我也觉得这些概念有点绕。但用多了就会发现，这套设计其实非常符合直觉——你只需要关心「单个样本」的计算逻辑，剩下的 batch 化和微分，JAX 帮你搞定。

嗯，今天就先聊到这里。记住一点：在自动驾驶感知中，正确的梯度 + 高效的向量化 = 可靠的模型训练。下次你写训练循环的时候，试试用 vmap 包一下 grad，你会回来感谢我的。

专注资料整理

← 上一章 📖 返回目录 下一章 →