深度学习框架实战对比：PyTorch、TensorFlow 还是 JAX？

坐标北京，通勤一小时，最近在研究 AI 相关技术，顺便记录一下踩过的坑。

---

开篇：为什么写这篇文章

上周五晚上九点半，我坐在工位上对着屏幕发呆。产品经理下午突然跑过来说："咱们那个内部知识库搜索体验太差了，能不能搞个 AI 搜索？"

我当时第一反应是：又来活了。

但转念一想，这不正好是我最近在学的方向吗？于是拍着胸脯说："没问题，给我两周。"

说完我就后悔了。

因为接下来的问题是：用什么框架？

PyTorch？TensorFlow？还是最近很火的 JAX？

作为一个在县城远程办公的小镇做题家（别问为什么坐标北京还说是县城，问就是精神故乡），我平时写博客记录学习心得，但这次是真的要上生产环境了，压力山大。

所以这篇文章，就是我在实际项目中对比这几个框架的真实体验。

---

背景：AI 搜索到底要做什么

先简单说一下需求。

公司有一个内部知识库，大概几十万篇文档，包括技术文档、产品手册、FAQ 等等。之前的搜索就是简单的关键词匹配，经常搜不到想要的东西。

现在要改成 AI 搜索，核心思路是：

1. 用 Embedding 模型把文档向量化
2. 用户提问时，把问题也向量化
3. 做向量相似度匹配，返回最相关的文档

听起来很简单对吧？

但实际做起来，坑是一个接一个。

---

框架选型：我的纠结过程

PyTorch：社区生态真的强

我第一个试的是 PyTorch。

原因很简单：Hugging Face 的 Transformers 库对 PyTorch 支持最好。

from transformers import AutoTokenizer, AutoModel
import torch

# 加载预训练模型
tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-large-zh")
model = AutoModel.from_pretrained("BAAI/bge-large-zh")

# 文本编码
def encode_text(texts):
    inputs = tokenizer(texts, padding=True, truncation=True, 
                       max_length=512, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    # 使用 [CLS] token 的表示作为句子向量
    embeddings = outputs.last_hidden_state[:, 0, :]
    return embeddings

优点：

• 调试方便，动态图真的很爽
• 社区资源丰富，遇到问题基本都能搜到解决方案
• Hugging Face 生态完善，模型下载一行代码搞定

缺点：

• 推理速度相对较慢（没有做专门优化的话）
• 部署到生产环境需要额外的工作

TensorFlow：工业级部署更成熟

然后我试了 TensorFlow。

说实话，TF 的代码写起来没有 PyTorch 那么直观，但它的优势在于部署。

import tensorflow as tf
import tensorflow_hub as hub

# 加载模型
model = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")

# 文本编码
def encode_text(texts):
    embeddings = model(texts)
    return embeddings

优点：

• TensorFlow Serving 部署真的很方便
• TFX 流水线适合大规模训练
• 移动端部署（TFLite）支持好

缺点：

• 调试体验不如 PyTorch
• 静态图虽然性能好，但写起来不够灵活
• 社区活跃度感觉在下降

JAX：性能怪兽但学习曲线陡

最后我试了 JAX。

JAX 是 Google 出的，主打函数式编程 + 自动微分 + XLA 编译优化。

import jax
import jax.numpy as jnp
from transformers import FlaxAutoModel, AutoTokenizer

# 加载模型
tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-large-zh")
model = FlaxAutoModel.from_pretrained("BAAI/bge-large-zh")

# JIT 编译加速
@jax.jit
def encode_text(input_ids, attention_mask):
    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
    return outputs.last_hidden_state[:, 0, :]

优点：

• 性能真的强，XLA 编译后速度起飞
• 函数式风格，代码可预测性好
• 支持多设备并行（pmap）

缺点：

• 学习曲线太陡了
• 调试困难，报错信息看不懂
• 生态还不够成熟

---

实战对比：用数据说话

光说优缺点没用，得上数据。

我在同样的硬件环境（一张 RTX 3090）上，对三个框架做了对比测试。

测试场景

• 模型：BAAI/bge-large-zh（中文 Embedding 模型）
• 数据集：10 万条文档，平均长度 200 字
• 任务：批量编码 + 相似度计算

性能数据

框架	编码速度（条/秒）	内存占用	部署难度	综合评分
PyTorch	850	4.2GB	中等	⭐⭐⭐⭐
TensorFlow	720	4.8GB	简单	⭐⭐⭐
JAX	1200	3.9GB	困难	⭐⭐⭐⭐⭐

从数据上看，JAX 的性能确实最强，但部署难度也是最高的。

---

踩坑记录：那些让我想砸电脑的瞬间

坑一：OOM（Out of Memory）

刚开始用 PyTorch 的时候，batch size 设了 64，直接 OOM。

RuntimeError: CUDA out of memory. Tried to allocate 2.00 GiB 
(GPU 0; 24.00 GiB total capacity; 22.50 GiB already allocated; 
0 bytes free; 22.80 GiB reserved in total by PyTorch)

当时真的想砸电脑。

后来发现是模型加载时没有用 torch.float16，改成半精度后内存占用直接减半。

model = AutoModel.from_pretrained(
    "BAAI/bge-large-zh",
    torch_dtype=torch.float16  # 关键：使用半精度
)

坑二：Fine-tuning 时的梯度爆炸

因为通用 Embedding 模型在我们领域的效果不够好，我决定做 Fine-tuning。

结果训练了 10 个 step，loss 直接变成 NaN。

Step 10: loss = nan, lr = 2e-5

排查了半天，发现是学习率太大了。Embedding 模型的 Fine-tuning 学习率要比普通分类任务小很多，一般用 1e-5 到 5e-6。

from transformers import get_linear_schedule_with_warmup

optimizer = torch.optim.AdamW(model.parameters(), lr=2e-6)  # 小学习率
scheduler = get_linear_schedule_with_warmup(
    optimizer, 
    num_warmup_steps=100,
    num_training_steps=10000
)

坑三：GitHub Copilot 的迷惑建议

说到 Fine-tuning，不得不吐槽一下 GitHub Copilot。

当时我在写训练循环，Copilot 自动补全了一段代码：

# Copilot 建议的代码
for batch in dataloader:
    outputs = model(**batch)
    loss = outputs.loss
    loss.backward()
    optimizer.step()

看起来很正常对吧？

但我忘了加 optimizer.zero_grad()，结果梯度一直累积，训练了 5 个 epoch 后模型直接崩了。

当时排查了两个小时才发现这个问题，真的想给 Copilot 差评。

---

最终选择：PyTorch + ONNX Runtime

经过一番折腾，我最终选择了 PyTorch 训练 + ONNX Runtime 推理 的方案。

原因：

1. PyTorch 训练体验最好，调试方便
2. 导出 ONNX 后，推理速度可以提升 2-3 倍
3. ONNX Runtime 部署简单，支持多种硬件

# 导出 ONNX 模型
import torch
from transformers import AutoModel

model = AutoModel.from_pretrained("BAAI/bge-large-zh")
model.eval()

dummy_input = {
    "input_ids": torch.randint(0, 1000, (1, 128)),
    "attention_mask": torch.ones(1, 128, dtype=torch.long)
}

torch.onnx.export(
    model,
    (dummy_input["input_ids"], dummy_input["attention_mask"]),
    "bge-large-zh.onnx",
    input_names=["input_ids", "attention_mask"],
    output_names=["last_hidden_state"],
    dynamic_axes={
        "input_ids": {0: "batch", 1: "sequence"},
        "attention_mask": {0: "batch", 1: "sequence"},
        "last_hidden_state": {0: "batch", 1: "sequence"}
    }
)

# 使用 ONNX Runtime 推理
import onnxruntime as ort

session = ort.InferenceSession("bge-large-zh.onnx")

def encode_text_onnx(texts):
    inputs = tokenizer(texts, padding=True, truncation=True, 
                       max_length=512, return_tensors="np")
    outputs = session.run(None, {
        "input_ids": inputs["input_ids"],
        "attention_mask": inputs["attention_mask"]
    })
    return outputs[0][:, 0, :]

---

AI 搜索效果：终于搞定了，开心

上线后效果还不错：

• 搜索准确率从 65% 提升到 89%
• 响应时间控制在 200ms 以内
• 用户反馈明显好转

产品经理终于不催我了，甚至请我喝了杯奶茶。

---

总结：几点心得体会

1. 框架选择要看场景：研究用 PyTorch，部署用 ONNX/TensorFlow，追求极致性能用 JAX
2. Fine-tuning 要注意细节：学习率、数据预处理、损失函数都很关键
3. 工具要用但别依赖：GitHub Copilot 能提高效率，但生成的代码一定要仔细检查
4. 性能优化是系统工程：模型量化、ONNX 导出、硬件加速，每一步都有讲究

最后，如果你也在做 AI 搜索相关的项目，欢迎交流。

毕竟，在县城远程办公的小镇做题家，也就这点爱好了。

---

本文写于 2024 年，北京，通勤地铁上