第2章环境搭建与初体验

第2章：环境搭建与初体验

说实话，每次带新人入门JAX，我最怕的就是环境搭建这一步。不是技术难，而是坑多。我自己当年装TPU版本时，折腾了整整一个下午，最后发现是驱动版本不对。嗯，这节课我们就来把这步走顺。

2.1 安装JAX：CPU/GPU/TPU版本

JAX的安装其实比PyTorch要清爽很多。它没有那么多乱七八糟的依赖，核心就一个命令。但不同硬件版本，你得选对。

CPU版本

如果你只是想先体验一下，或者手头没有NVIDIA显卡，那就装CPU版。我个人习惯先用这个版本做原型验证。

pip install --upgrade pip
pip install jax jaxlib

注意：CPU版虽然能跑，但速度嘛...你想想看，就像开着一辆奥拓去跑赛道。能跑，但别指望推背感。

GPU版本（CUDA 11/12）

这是大多数人的选择。我建议先确认你的CUDA版本：

nvidia-smi  # 查看驱动版本和CUDA版本

然后根据CUDA版本选择安装命令。以CUDA 12为例：

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

这里有个坑：我曾经遇到过pip自动安装了CPU版，明明GPU驱动好好的。后来发现是pip源的问题。建议加上-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html这个参数。

TPU版本

TPU版本比较特殊。你需要在Google Cloud上创建TPU VM，或者用Colab。我个人建议初学者先别碰TPU，除非你真的需要大规模训练。

# 在TPU VM上
pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

⚠️ 避坑指南
我曾经在Windows上折腾JAX GPU版，结果发现官方根本不支持Windows原生安装。解决办法是用WSL2。如果你非要在Windows上跑，建议直接用Docker镜像。

2.2 JAX与NumPy的第一次亲密接触

JAX最迷人的地方，就是它和NumPy的API几乎一模一样。说白了，你只要会NumPy，JAX就上手了一半。

来看个最简单的例子：

import jax.numpy as jnp
import numpy as np

# 创建数组
x_np = np.array([1, 2, 3, 4, 5])
x_jax = jnp.array([1, 2, 3, 4, 5])

print("NumPy数组:", x_np)
print("JAX数组:", x_jax)
print("类型:", type(x_jax))

输出结果：

NumPy数组: [1 2 3 4 5]
JAX数组: [1 2 3 4 5]
类型: <class 'jaxlib.xla_extension.ArrayImpl'>

看到没？数据一样，但类型不同。JAX的数组是ArrayImpl，不是NumPy的ndarray。为什么？因为JAX的数组是设备数组，它可能躺在GPU显存里，而不是CPU内存里。

再试试运算：

# 矩阵乘法
a = jnp.array([[1, 2], [3, 4]])
b = jnp.array([[5, 6], [7, 8]])
c = jnp.dot(a, b)
print(c)

嗯，这里要注意：JAX默认是单精度（float32）的。如果你用NumPy的习惯写float64，JAX会报错。我刚开始就踩过这个坑，调试了半天才发现是精度问题。

💡 小技巧
如果你需要双精度，可以在代码开头加上：jax.config.update("jax_enable_x64", True)。但注意，GPU上双精度速度会慢很多。

2.3 JAX设备数组基础

设备数组，说白了就是放在GPU或TPU上的数组。这是JAX和NumPy最本质的区别。

查看设备信息

import jax

# 查看可用设备
print("可用设备:", jax.devices())
print("设备数量:", jax.device_count())
print("默认设备:", jax.default_device())

如果你有GPU，输出会类似这样：

可用设备: [gpu:0, gpu:1]
设备数量: 2
默认设备: gpu:0

显式控制设备

JAX允许你手动指定数组放在哪个设备上。我个人习惯在调试时用这个功能：

from jax import device_put

# 创建一个NumPy数组
x = np.array([1, 2, 3])

# 显式放到GPU上
x_gpu = device_put(x, device=jax.devices()[0])
print("设备:", x_gpu.device())

设备间的数据移动

这里有个重要概念：JAX的数组默认是不可变的。你不能像NumPy那样直接修改数组元素。为什么？因为JAX要保证函数式编程的纯度，这样才能做自动微分和JIT编译。

# 这行会报错！
# x_gpu[0] = 10  

# 正确做法：创建新数组
x_new = x_gpu.at[0].set(10)
print("原数组:", x_gpu)
print("新数组:", x_new)

🔑 核心要点

JAX的数组操作不会修改原数组，而是返回一个新数组。这跟NumPy的in-place操作完全不同。刚开始可能会不习惯，但习惯了你会发现这其实更安全。

2.4 知识体系图

下面这张图总结了本章的核心逻辑：

2.5 实战小练习

光看不练假把式。来，我们写个完整的例子，把今天学的串起来：

import jax
import jax.numpy as jnp
import numpy as np
import time

# 1. 检查环境
print("JAX版本:", jax.__version__)
print("可用设备:", jax.devices())

# 2. 创建数据
size = 1000
x_np = np.random.randn(size, size).astype(np.float32)
y_np = np.random.randn(size, size).astype(np.float32)

# 3. 放到GPU上
x_jax = jnp.array(x_np)
y_jax = jnp.array(y_np)

# 4. 矩阵乘法对比
# NumPy
start = time.time()
z_np = np.dot(x_np, y_np)
np_time = time.time() - start

# JAX
start = time.time()
z_jax = jnp.dot(x_jax, y_jax)
jax_time = time.time() - start

print(f"NumPy耗时: {np_time:.4f}秒")
print(f"JAX耗时: {jax_time:.4f}秒")
print(f"加速比: {np_time/jax_time:.2f}x")

运行这个代码，你会看到JAX比NumPy快很多。我第一次跑的时候，看到加速比10x以上，说实话挺震撼的。这就是硬件加速的魅力。

💡 个人经验
如果你发现加速比不大，别急着怀疑。检查一下：1) 数据量够不够大？2) 是不是第一次运行（JIT编译需要预热）？3) 数组是不是float32？我遇到过有人用float64，结果GPU跑得比CPU还慢。

好了，环境搭好了，代码也跑了。下一节我们会深入JAX的核心机制——JIT编译和自动微分。嗯，那才是真正好玩的地方。

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

📚 课程目录