深夜在成都用AI编程工具跑深度学习框架的那些事

凌晨两点半，成都的夜已经深了。窗外偶尔传来几声出租车的喇叭，楼下那家串串香早就关门了。我灌了第三罐可乐，盯着屏幕上PyTorch和TensorFlow的训练loss曲线，顺手在Trae里敲下了一行注释。

说实话，写这篇文章的契机挺有意思的。上个月我们组接了个计算机视觉的项目，要给一家本地餐饮连锁做菜品识别系统。说白了就是让摄像头拍一下盘子，自动算出你吃了啥、多少钱。产品经理说这个需求"很简单"的时候，我差点把手里的盖碗茶泼他脸上。

作为一个白天写业务CRUD、晚上准备考公的打工人，我本来是想用这项目练练手，顺便把最近研究的Rust和深度学习结合一下。结果一上来就面临一个灵魂拷问：用哪个框架？

框架选型这件事，真没想的那么简单

我们团队之前主要用TensorFlow，毕竟历史包袱在那儿，老系统都是TF 1.x时代的产物。但说实话，TF 2.x虽然改进了很多，动态图用起来也顺了，可社区活跃度肉眼可见地在下滑。Stack Overflow上随便搜个问题，排在前面的回答很多都是2019年以前的。

PyTorch这边呢，学术圈基本全面倒戈了。我翻了几篇最新的论文，代码清一色PyTorch。Hugging Face的生态也是围着PyTorch转的。但问题是，我们线上部署环境是TensorFlow Serving，切换框架意味着整个推理链路都要重构。

JAX我也看了，Google亲儿子，函数式编程的风格确实很优雅。但团队里没几个人熟，学习成本太高。而且JAX的调试体验……懂的都懂，一个jit编译错误能让你怀疑人生。

就在我纠结的时候，领导甩了句话："你先做个技术调研，下周给我结论。"

得，又是个周末要搭进去的活儿。

借助AI工具加速调研过程

既然要对比，那就得真刀真枪地跑。我给自己定了个计划：用同一个数据集、同一个模型结构，分别在三个框架上实现，然后对比开发效率、训练速度、部署便利性。

数据集用的是Food-101，101个菜品分类，每个分类1000张图。模型选了ResNet-50，不算太复杂，但足够说明问题。

这里要说说我最近用的几个AI辅助工具，确实帮了大忙。

先说Trae，这是字节跳动出的IDE，基于VS Code魔改的。我一开始是冲着它的免费额度去的（穷逼程序员的基操），结果用下来发现体验还真不错。它内置的AI对话功能可以直接在编辑器里问问题，不用来回切窗口。我写PyTorch的DataLoader时，有个多进程预取的参数调了半天不对，直接在Trae里问，它给我分析了num_workers和pin_memory的配合问题，还顺带提醒了我persistent_workers在较新版本里的变化。省了我至少半小时查文档的时间。

Tabnine我也在用，主要是代码补全。这玩意儿训练数据据说比较干净，补全的代码质量相对稳定。写TensorFlow的tf.data pipeline时，它的补全确实能跟上思路，特别是那些又臭又长的API名字，tf.data.experimental.parallel_interleave这种，手打真的会死。不过它有个问题，对中文注释的理解偶尔会抽风，我写个"这里要做数据增强"，它有时候会给我补一段完全不相关的代码。

GitHub Copilot是团队统一采购的，说实话，在深度学习代码这块，它的表现有点两极分化。写一些标准的训练循环、评估逻辑，它补全得很准，基本上我写个def train_epoch(，它就能把整个函数骨架给你填好。但涉及到一些框架特有的骚操作，比如PyTorch的torch.compile或者TensorFlow的自定义训练步骤，它经常给你一些看起来对但实际跑不通的代码。有一次它给我补了段TF的梯度裁剪代码，用了个已经deprecated的API，我debug了二十分钟才发现。

Perplexity是我用来做技术调研的主力。这玩意儿本质上是个AI搜索引擎，但比传统的搜索好用太多了。我想知道"PyTorch 2.0的torch.compile在ResNet上的实际加速比"，直接问它，它会给我整理好相关的benchmark数据、论文引用、社区讨论，还会标注信息来源。比起自己在Google上翻十几个网页，效率高了一个量级。

工具	主要用途	优点	槽点
Trae	代码编写+AI对话	免费、集成度高、中文支持好	偶尔响应慢
Tabnine	代码补全	补全质量稳定、本地模型可选	中文注释理解偶尔抽风
GitHub Copilot	代码补全+生成	团队统一采购、长代码生成强	框架新特性支持滞后
Perplexity	技术调研	信息来源可追溯、结构化回答好	深度技术问题分析偶尔不准

三个框架的实际体验

PyTorch：写起来真爽

先说PyTorch。代码写起来确实舒服，动态图的调试体验没得说。我写模型的时候，直接print中间tensor的shape，断点调试想看哪里看哪里，不用像TF 1.x那样搞一堆tf.print。

import torch
import torch.nn as nn
from torchvision import models, transforms
from torch.utils.data import DataLoader

# 数据预处理
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                         std=[0.229, 0.224, 0.225])
])

# 模型定义，继承nn.Module就完事了
class FoodClassifier(nn.Module):
    def __init__(self, num_classes=101):
        super().__init__()
        self.backbone = models.resnet50(weights='IMAGENET1K_V2')
        self.backbone.fc = nn.Linear(2048, num_classes)
    
    def forward(self, x):
        return self.backbone(x)

model = FoodClassifier()
# 看看模型结构，直接print就行
dummy_input = torch.randn(1, 3, 224, 224)
output = model(dummy_input)
print(f"Output shape: {output.shape}")  # 直接看，不用跑session

训练循环也很直观，没有什么黑盒：

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss()

for epoch in range(20):
    model.train()
    running_loss = 0.0
    for batch_idx, (images, labels) in enumerate(train_loader):
        images, labels = images.to(device), labels.to(device)
        
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
        if batch_idx % 50 == 0:
            print(f"Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item():.4f}")

PyTorch 2.0的torch.compile我也试了，加了就一行代码：

model = torch.compile(model, mode="reduce-overhead")

在Food-101上，训练速度大概提升了15%-20%，没吹的那么神，但白嫖的性能不要白不要。

不过PyTorch的部署确实是个痛点。虽然TorchServe能用，但跟TensorFlow Serving比，稳定性和生态还是差一截。我们运维老哥看到TorchServe的文档时，那个表情我至今难忘。

TensorFlow：老当益壮但略显疲态

TF的代码写起来确实啰嗦一些，特别是Keras API和底层API混用的时候。但不得不说，TF的部署生态是真的成熟。TF Serving、TF Lite、TF.js，全家桶安排得明明白白。

import tensorflow as tf
from tensorflow import keras

# 数据pipeline，tf.data确实强大
def preprocess(image, label):
    image = tf.image.resize(image, [256, 256])
    image = tf.image.central_crop(image, 224/256)
    image = tf.keras.applications.resnet.preprocess_input(image)
    return image, label

train_ds = tf.data.Dataset.from_tensor_slices((train_images, train_labels))
train_ds = train_ds.map(preprocess, num_parallel_calls=tf.data.AUTOTUNE)
train_ds = train_ds.batch(32).prefetch(tf.data.AUTOTUNE)

# 模型定义，Keras用起来还是方便的
base_model = keras.applications.ResNet50V2(
    weights='imagenet', 
    include_top=False,
    input_shape=(224, 224, 3)
)
base_model.trainable = False  # 先冻结

inputs = keras.Input(shape=(224, 224, 3))
x = base_model(inputs, training=False)
x = keras.layers.GlobalAveragePooling2D()(x)
outputs = keras.layers.Dense(101, activation='softmax')(x)
model = keras.Model(inputs, outputs)

model.compile(
    optimizer=keras.optimizers.Adam(1e-4),
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# 训练，一行搞定
history = model.fit(
    train_ds,
    epochs=20,
    validation_data=val_ds,
    callbacks=[
        keras.callbacks.EarlyStopping(patience=3),
        keras.callbacks.ReduceLROnPlateau(factor=0.5)
    ]
)

TF的tf.data API确实强，数据管道的性能优化做得很好。但写起来那个链式调用，有时候看得人眼花。而且Keras的高层API和底层TF操作混在一起时，调试起来很痛苦。你用了个自定义的tf.function，结果梯度传不过去，报错信息又写得跟天书一样，那种感觉真的很崩溃。

上周五晚上我就遇到了这个问题。一个自定义的loss function，在eager模式下跑得好好的，一@tf.function就nan。折腾到凌晨一点，最后发现是tf.where的条件分支里有个梯度断裂的问题。当时真的想砸键盘。

JAX：优雅但劝退

JAX我只做了个简单的demo，没敢上生产。函数式编程的风格确实很纯粹，但心智负担太重了。

import jax
import jax.numpy as jnp
from flax import linen as nn

class ResNetBlock(nn.Module):
    features: int
    strides: tuple = (1, 1)
    
    @nn.compact
    def __call__(self, x, train=True):
        residual = x
        y = nn.Conv(self.features, (3, 3), self.strides, padding='SAME')(x)
        y = nn.BatchNorm(use_running_average=not train)(y)
        y = nn.relu(y)
        y = nn.Conv(self.features, (3, 3), (1, 1), padding='SAME')(y)
        y = nn.BatchNorm(use_running_average=not train)(y)
        
        if self.strides != (1, 1):
            residual = nn.Conv(self.features, (1, 1), self.strides)(x)
        
        return nn.relu(y + residual)

# 初始化和调用，一切皆函数
model = ResNetBlock(features=64)
key = jax.random.PRNGKey(0)
params = model.init(key, jnp.ones((1, 224, 224, 3)))

# JIT编译，性能拉满
@jax.jit
def train_step(params, batch, key):
    def loss_fn(params):
        logits = model.apply(params, batch['image'], train=True, rngs={'dropout': key})
        return jnp.mean(-jnp.sum(batch['label'] * jax.nn.log_softmax(logits), axis=-1))
    
    loss, grads = jax.value_and_grad(loss_fn)(params)
    return loss, grads

JAX的jit编译确实猛，在同样的硬件上，纯计算速度比PyTorch快不少。但问题是，每次修改代码都要重新编译，调试的时候那个等待时间真的让人抓狂。而且vmap、pmap这些批量化的操作，理解起来有门槛。

对于我们这种要快速迭代的项目，JAX的开发效率还是不够。除非你是做科研的，追求极致的计算性能，否则不建议轻易入坑。

对比总结

跑了两周的benchmark，数据汇总如下：

指标	PyTorch 2.0	TensorFlow 2.15	JAX 0.4
开发效率	⭐⭐⭐⭐⭐	⭐⭐⭐	⭐⭐
训练速度(单卡)	1x	0.92x	1.15x
调试体验	⭐⭐⭐⭐⭐	⭐⭐⭐	⭐⭐
部署便利性	⭐⭐⭐	⭐⭐⭐⭐⭐	⭐⭐
社区活跃度	⭐⭐⭐⭐⭐	⭐⭐⭐	⭐⭐⭐
学习曲线	平缓	中等	陡峭

最后给领导的结论是：新项目上PyTorch，老系统维护继续TF，JAX暂时观望。

领导问我："那部署怎么办？"

我说："TorchServe我研究过了，文档虽然烂点，但能用。实在不行就ONNX导出来，TensorRT推理也行。"

领导点点头："行，那你负责把TorchServe的部署流程跑通，写个文档给运维。"

我："……"

得，又给自己挖了个坑。

考公和写代码，其实没那么冲突

说到这儿，可能有人好奇，我都准备考公了，还折腾这些干嘛？

其实吧，考公是给自己留条后路。程序员这行，35岁的焦虑不是假的。但在职的这段时间，该学的还是得学，该做的項目还是得做好。一方面是为了对得起这份工资，另一方面，这些实战经验不管以后干啥都用得上。

而且说实话，深度学习这块的东西，考公面试有时候也会问到。特别是现在各地都在搞数字政府、智慧城市，懂点AI的技术岗还是吃香的。

成都这边的生活节奏确实舒服，不像北京上海那么卷。晚上下班回家，吃个火锅，写写代码，看看书，日子过得还行。深夜写代码效率确实高，可能是没人打扰的缘故。

最近研究Rust也是类似的想法。Rust的所有权机制、内存安全这些概念，跟考公复习的逻辑推理其实有异曲同工之妙。都是要你把规则吃透，然后在规则内找到最优解。

好了，可乐喝完了，代码也跑完了。训练loss已经降到0.3了，明天看看验证集的效果。

如果这篇文章对你有帮助，点个赞呗。如果没帮助……那你来成都，我请你吃串串。