2、环境搭建与项目初始化：安装JAX与Flax、配置GPU/TPU环境、创建第一个Flax项目、理解项目结构

说实话，很多同学学深度学习框架，上来就急着写模型。结果环境没配好，跑个demo报一堆错，心态直接崩了。我当年刚接触JAX时也踩过不少坑，尤其是GPU版本死活装不上，折腾了一整天。

这一章，咱们就把地基打牢。环境搭好了，后面才能跑得顺畅。

2.1 安装JAX：CPU、GPU、TPU版本怎么选？

JAX的安装其实很简单，但有个关键点——你得先搞清楚自己用的是什么硬件。

核心原则：JAX的安装包是分硬件平台的。CPU版本和GPU版本是两套不同的wheel包。

CPU版本（适合学习调试）

如果你只是想先跑跑代码，或者电脑没有NVIDIA显卡，直接装CPU版就行：

pip install --upgrade pip
pip install jax jaxlib

嗯，就这么简单。装完就能用了，但跑大模型会慢一些。

GPU版本（生产环境首选）

我个人习惯用GPU版本做训练。前提是你得有NVIDIA显卡，并且装好了CUDA和cuDNN。

先确认CUDA版本：

nvcc --version

然后根据CUDA版本选择对应的jaxlib。以CUDA 12为例：

pip install --upgrade pip
pip install jax[cuda12]

⚠️ 我曾经踩过的坑：直接装jax[cuda12]时，如果pip源里没有对应版本，它会自动降级装CPU版。跑代码时不报错，但速度慢得离谱。建议装完后用 jax.devices() 检查一下。

TPU版本（Google云专用）

如果你用的是Google Cloud TPU，那需要装TPU版：

pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

说实话，TPU版本我平时用得不多，主要是成本高。但如果你做大规模分布式训练，TPU确实香。

2.2 配置GPU/TPU环境：验证你的硬件是否就绪

装完JAX后，第一件事就是验证环境。别急着写模型，先跑几行代码确认硬件被正确识别。

检查可用设备

import jax
print(jax.devices())

如果输出类似 [gpu(id=0)] 或 [tpu(id=0)]，说明环境配好了。如果输出 [cpu(id=0)]，嗯，你装的是CPU版。

测试GPU是否真的在工作

我有个小习惯，会跑一个简单的矩阵乘法来验证：

import jax.numpy as jnp
import time

# 创建一个大矩阵
a = jnp.ones((5000, 5000))
b = jnp.ones((5000, 5000))

start = time.time()
c = jnp.dot(a, b)
jax.block_until_ready(c)  # 强制等待计算完成
print(f"计算耗时: {time.time() - start:.4f}秒")

GPU版本通常比CPU快几十倍。如果速度没差别，那可能是JAX没用到GPU。

💡 避坑指南：JAX默认是惰性计算（lazy evaluation），你写 c = jnp.dot(a, b) 时它并不会立刻算。必须用 jax.block_until_ready() 或 .block_until_ready() 才能拿到实际结果。我刚开始学JAX时就被这个坑过，以为代码写错了。

2.3 安装Flax：轻量级神经网络库

Flax是JAX生态里最流行的神经网络库。它不像PyTorch那样封装得特别重，而是给你更多底层控制权。

安装很简单：

pip install flax

装完后验证一下：

import flax
print(flax.__version__)

如果没报错，恭喜你，环境搭好了。

2.4 创建第一个Flax项目

我个人习惯用模块化的方式组织项目。别把所有代码塞进一个文件里，那样后期维护会想哭。

先创建项目目录：

mkdir my_first_flax_project
cd my_first_flax_project

然后创建以下文件结构：

my_first_flax_project/
├── main.py          # 入口文件
├── model.py         # 模型定义
├── train.py         # 训练逻辑
├── data.py          # 数据加载
├── config.py        # 超参数配置
└── requirements.txt # 依赖清单

你想想看，这样分文件的好处是什么？

  • model.py：只放模型结构，改模型时不用翻其他代码
  • train.py：只放训练循环，调参时一目了然
  • config.py：所有超参数集中管理，方便实验对比

写一个最简单的模型

在 model.py 里定义一个线性回归模型：

import flax.linen as nn

class LinearRegression(nn.Module):
    @nn.compact
    def __call__(self, x):
        return nn.Dense(features=1)(x)

然后在 main.py 里调用：

import jax
import jax.numpy as jnp
from model import LinearRegression

# 初始化模型
model = LinearRegression()
params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 10)))

# 前向传播
x = jnp.ones((5, 10))
output = model.apply(params, x)
print(output)

跑一下看看，如果输出一个5x1的矩阵，说明项目跑通了。

注意：Flax的模型初始化需要传入一个伪随机数种子（PRNGKey），这是JAX的纯函数设计决定的。每次初始化时，如果你用同一个key，得到的结果完全一样。这在调试时特别有用。

2.5 理解项目结构：Flax项目的核心逻辑

Flax项目的结构其实有规律可循。说白了，就三个核心概念：

概念 作用 对应文件
Module 定义神经网络层和参数 model.py
State 管理参数、优化器状态等 train.py
Transform 实现梯度计算、编译优化等 train.py

我画了一张图，帮你理清它们之间的关系：

Flax项目核心结构 Module 定义网络层 声明参数结构 State 管理参数值 优化器状态 Transform 梯度计算 JIT编译 典型数据流 1. 定义Module → 描述网络结构（如LinearRegression） 2. 初始化State → 生成随机参数（model.init） 3. 应用Transform → 计算梯度、更新参数（jax.grad, jit） 4. 前向传播 → 用model.apply执行推理

你看，这个流程其实很清晰。Module负责「长什么样」，State负责「当前值是多少」，Transform负责「怎么更新」。三者各司其职。

2.6 实战：跑通你的第一个Flax训练脚本

光说不练假把式。咱们写一个完整的训练脚本，用随机数据训练线性回归模型。

在 train.py 里添加：

import jax
import jax.numpy as jnp
import optax  # Flax官方推荐的优化器库
from model import LinearRegression

def train():
    # 生成随机数据
    key = jax.random.PRNGKey(42)
    x = jax.random.normal(key, (100, 10))
    y = jnp.sum(x, axis=1, keepdims=True)  # 目标：所有特征之和
    
    # 初始化模型和优化器
    model = LinearRegression()
    params = model.init(key, x[:1])
    optimizer = optax.sgd(learning_rate=0.01)
    opt_state = optimizer.init(params)
    
    # 定义损失函数
    def loss_fn(params, x, y):
        pred = model.apply(params, x)
        return jnp.mean((pred - y) ** 2)
    
    # 编译训练步骤
    @jax.jit
    def train_step(params, opt_state, x, y):
        loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
        updates, opt_state = optimizer.update(grads, opt_state, params)
        params = optax.apply_updates(params, updates)
        return params, opt_state, loss
    
    # 训练循环
    for epoch in range(100):
        params, opt_state, loss = train_step(params, opt_state, x, y)
        if epoch % 10 == 0:
            print(f"Epoch {epoch}, Loss: {loss:.6f}")

if __name__ == "__main__":
    train()

跑一下：

python train.py

如果看到loss在下降，说明一切正常。嗯，就是这么简单。

💡 我的经验：第一次跑通Flax项目时，我建议你用 @jax.jit 装饰训练步骤。它会把Python代码编译成高效的XLA计算图，速度能提升10倍以上。但注意，jit对代码有要求——不能有Python控制流（if/for）依赖动态数据。如果报错，先去掉jit调试。

2.7 本章小结

这一章我们做了几件事：

  • 安装了JAX和Flax，区分了CPU/GPU/TPU版本
  • 验证了硬件环境是否正常工作
  • 创建了第一个Flax项目，理解了模块化结构
  • 跑通了一个完整的训练脚本

说实话，环境搭建是最枯燥的部分，但也是最重要的。我见过太多同学因为环境问题卡在第一步，最后放弃了整个框架。所以，花点时间把这一步走稳，后面你会感谢自己的。

专注资料整理

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