1. JAX基础与环境搭建

各位同学好，欢迎来到JAX图像处理与生成模型实现的第一章。

说实话，我第一次接触JAX是在2020年。当时我在做一个生成对抗网络的项目，TensorFlow和PyTorch都试过了，总觉得梯度计算和性能优化上差点意思。后来一个同事跟我说：「试试JAX吧，这玩意儿写起来像NumPy，跑起来像飞一样。」嗯，一试就回不去了。

1.1 JAX是什么？

JAX，全称是「Just Another XLA」，但我觉得它一点也不「Just Another」。它是Google开发的一个Python库，专门做高性能数值计算和机器学习研究。

你想想看，它把NumPy的API、自动微分、GPU/TPU加速、即时编译（JIT）全揉在一起了。说白了，就是让你用写NumPy的语法，写出能自动求导、能并行计算、能编译优化的代码。

核心定位：JAX = NumPy + 自动微分 + GPU/TPU加速 + 即时编译

1.2 安装配置

安装JAX其实很简单。我个人习惯用pip直接装：

# CPU版本（适合学习和调试）
pip install jax jaxlib

# GPU版本（需要CUDA和cuDNN）
pip install jax[cuda12] jaxlib

注意：GPU版本对CUDA版本有要求。我曾经在项目里踩过坑——装了最新的CUDA 12.2，结果jaxlib不兼容，折腾了半天。建议先查一下官方文档的版本对应表。

验证安装是否成功：

import jax
import jax.numpy as jnp

print(jax.__version__)
print(jax.devices())  # 看看有没有检测到GPU

如果输出里能看到 [cuda:0] 之类的信息，那就恭喜你，GPU加速已经就绪了。

1.3 JAX与NumPy的对比

先看个直观的例子：

import numpy as np
import jax.numpy as jnp

# NumPy写法
x_np = np.array([1.0, 2.0, 3.0])
y_np = np.sin(x_np) + np.cos(x_np)

# JAX写法（几乎一模一样）
x_jax = jnp.array([1.0, 2.0, 3.0])
y_jax = jnp.sin(x_jax) + jnp.cos(x_jax)

print(y_np)   # [1.3818 1.3254 0.8489]
print(y_jax)  # [1.3818 1.3254 0.8489]

看到没？API几乎一模一样。但有几个关键区别：

特性 NumPy JAX
数组可变性 支持原地修改 数组不可变（函数式）
自动微分 不支持 原生支持（grad）
GPU加速 不支持 原生支持
即时编译 不支持 支持（jit）
随机数生成 有状态 无状态（需要显式传入key）

小技巧：刚开始用JAX时，最容易犯的错误就是试图修改数组。比如 x[0] = 1.0 在NumPy里没问题，但在JAX里会报错。记住：JAX的数组是不可变的，你需要用 x.at[0].set(1.0) 这种函数式写法。

1.4 JAX的核心概念

JAX有三个核心概念，我管它们叫「三驾马车」：jit、vmap、grad。搞懂了这三个，你就掌握了JAX的精髓。

1.4.1 jit：即时编译

jit，全称Just-In-Time compilation。说白了，就是把你的Python函数编译成高效的XLA代码，跑起来快得飞起。

import jax
import jax.numpy as jnp
import time

def slow_func(x):
    for _ in range(1000):
        x = jnp.sin(x) + jnp.cos(x)
    return x

@jax.jit
def fast_func(x):
    for _ in range(1000):
        x = jnp.sin(x) + jnp.cos(x)
    return x

x = jnp.ones((1000,))

# 第一次调用会编译，稍慢
start = time.time()
result_slow = slow_func(x)
print(f"未jit: {time.time() - start:.4f}秒")

# jit版本第一次调用也会编译
start = time.time()
result_fast = fast_func(x)
print(f"jit第一次: {time.time() - start:.4f}秒")

# 第二次调用就快了
start = time.time()
result_fast = fast_func(x)
print(f"jit第二次: {time.time() - start:.4f}秒")

经验之谈：我在项目中遇到过一个问题——jit编译后的函数不能包含Python控制流（比如if语句依赖于数组的值）。因为jit在编译时就需要知道所有可能的执行路径。解决办法是用 jax.lax.cond 或 jax.lax.while_loop 替代。

1.4.2 vmap：自动向量化

vmap，全称vectorizing map。它的作用是把一个处理单个样本的函数，自动变成能处理批量样本的函数。

import jax
import jax.numpy as jnp

# 处理单个样本的函数
def process_single(x):
    return jnp.sin(x) + jnp.cos(x)

# 批量处理：手动写循环
def process_batch_manual(batch):
    results = []
    for i in range(batch.shape[0]):
        results.append(process_single(batch[i]))
    return jnp.stack(results)

# 批量处理：用vmap自动向量化
process_batch_vmap = jax.vmap(process_single)

# 测试
batch = jnp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

result_manual = process_batch_manual(batch)
result_vmap = process_batch_vmap(batch)

print(jnp.allclose(result_manual, result_vmap))  # True

避坑指南：我曾经在图像处理项目里，用vmap把单张图片的处理函数向量化到整个batch。结果发现vmap自动处理了batch维度，但我的函数里还有通道维度...嗯，后来用 in_axes 和 out_axes 参数指定了维度映射才搞定。

1.4.3 grad：自动微分

grad，就是自动求梯度。这是深度学习的基础，JAX的grad用起来特别顺手。

import jax
import jax.numpy as jnp

# 定义一个简单的函数
def f(x):
    return x**2 + 3*x + 1

# 求一阶导数
df = jax.grad(f)
print(df(2.0))  # 2*2 + 3 = 7.0

# 求二阶导数
d2f = jax.grad(jax.grad(f))
print(d2f(2.0))  # 2.0

# 多参数求导
def g(x, y):
    return x**2 + y**3

# 只对第一个参数求导
dg_dx = jax.grad(g, argnums=0)
print(dg_dx(2.0, 3.0))  # 2*2 = 4.0

# 对两个参数都求导
dg_dx_dy = jax.grad(g, argnums=(0, 1))
print(dg_dx_dy(2.0, 3.0))  # (4.0, 27.0)

核心要点：grad返回的是一个函数，不是数值。这个函数接受和原函数一样的参数，返回梯度。你想想看，这意味着你可以把grad和jit、vmap组合使用——比如 jax.jit(jax.vmap(jax.grad(f)))，这就是JAX的威力所在。

1.5 本章小结

这一章我们走了一遍JAX的基础：

  • JAX是什么——NumPy的升级版，自带自动微分和GPU加速
  • 怎么安装——pip搞定，注意CUDA版本兼容
  • 和NumPy的区别——数组不可变、无状态随机数、函数式编程风格
  • 三个核心概念——jit（编译加速）、vmap（自动向量化）、grad（自动求导）

说实话，JAX的学习曲线比PyTorch稍微陡一点，但一旦你习惯了函数式编程的思维，你会发现它写起来特别干净、特别优雅。下一章我们会开始用JAX做实际的图像处理操作，到时候你会更深刻地体会到这些概念的价值。

公众号:蓝海资料掘金营，微信deep3321

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