第三章 深度学习框架优化：PyTorch性能调优、TensorFlow性能调优、框架对比与选择
说实话，框架优化这事儿，我踩过的坑比走过的路还多。刚开始带团队的时候，大家总觉得模型跑得慢是显卡不够好，后来我发现——很多时候，是框架没用好。你想想看，同样的模型，有人用PyTorch跑得飞起，有人用TensorFlow卡成PPT，差别在哪？就在这些调优细节里。
3.1 PyTorch性能调优：从“能用”到“好用”
PyTorch是我个人最常用的框架，没有之一。它的动态图机制让调试变得特别舒服，但舒服归舒服，性能上稍不注意就会翻车。
3.1.1 DataLoader的“隐形杀手”
我在项目中遇到过最离谱的一次——训练速度只有理论值的30%。排查了半天，发现是DataLoader的num_workers设成了0。说白了，数据加载成了单线程，GPU在那儿干等着。
核心要点：
num_workers一般设为CPU核心数的2-4倍，但别超过8，否则调度开销反而大- 记得加
pin_memory=True，尤其是用GPU训练时，能省掉CPU到GPU的数据拷贝时间 - 如果数据预处理很重，考虑用
prefetch_factor预取更多批次
# 我常用的DataLoader配置
from torch.utils.data import DataLoader
dataloader = DataLoader(
dataset,
batch_size=64,
shuffle=True,
num_workers=4, # 根据CPU核心数调整
pin_memory=True, # 这个一定要开
prefetch_factor=2 # 预取因子
)
3.1.2 梯度累积与混合精度
显存不够怎么办？很多人第一反应是换小卡。其实不用，梯度累积就能解决。我曾经在8GB显存的卡上训过BERT-large，靠的就是这招。
我的经验：梯度累积步数一般设2-8步。设太大，模型收敛会变慢，因为参数更新频率降低了。
# 梯度累积示例
accumulation_steps = 4
optimizer.zero_grad()
for i, (inputs, labels) in enumerate(dataloader):
outputs = model(inputs)
loss = criterion(outputs, labels)
loss = loss / accumulation_steps # 归一化
loss.backward()
if (i + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
混合精度训练（AMP）就更不用说了。PyTorch 1.6以后原生支持，我建议所有训练任务都开。为什么？因为FP16能把训练速度提升1.5-2倍，显存占用还减半。
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for inputs, labels in dataloader:
optimizer.zero_grad()
with autocast():
outputs = model(inputs)
loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
注意：AMP不是万能的。有些操作（比如某些自定义loss）在FP16下会溢出。我曾经遇到过，训练到一半loss变成NaN，排查了半天才发现是某个除法操作精度不够。解决办法：把那个操作放到with autocast(enabled=False):里跑。
3.2 TensorFlow性能调优：静态图的威力
TensorFlow 2.x虽然默认开启动态图（Eager Execution），但真正发挥性能还得靠tf.function。说白了，就是把Python代码编译成静态计算图，跑起来快得多。
3.2.1 tf.function的正确用法
我记得刚转TF2的时候，觉得@tf.function就是个装饰器，加就完事了。结果发现，加错了地方反而更慢。为什么？因为tf.function会把Python控制流（if/for）转换成TF的图操作，如果转换失败，它会回退到Eager模式，性能反而下降。
最佳实践：
- 把
@tf.function加在训练循环的外层，而不是内层 - 避免在
tf.function内部使用Python的print、list.append等操作 - 用
tf.print代替print，用tf.TensorArray代替list
@tf.function
def train_step(inputs, labels):
with tf.GradientTape() as tape:
predictions = model(inputs, training=True)
loss = loss_fn(labels, predictions)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return loss
# 训练循环
for epoch in range(epochs):
for inputs, labels in dataset:
loss = train_step(inputs, labels) # 这里调用的是编译后的图
3.2.2 数据管道的优化
TensorFlow的tf.data API非常强大，但很多人只用了个皮毛。我见过最夸张的案例——有人用tf.data.Dataset.from_tensor_slices加载整个数据集到内存，结果16GB内存直接爆了。
正确的做法是用tf.data的流水线优化：
dataset = tf.data.Dataset.list_files('data/*.tfrecord')
dataset = dataset.interleave(
lambda x: tf.data.TFRecordDataset(x),
cycle_length=4, # 并行读取的文件数
num_parallel_calls=tf.data.AUTOTUNE
)
dataset = dataset.map(parse_function, num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.batch(64).prefetch(tf.data.AUTOTUNE)
避坑指南：我曾经把prefetch放在batch前面，结果每个样本都单独预取，效率极低。记住：prefetch一定要放在最后，而且参数用tf.data.AUTOTUNE让框架自动调整。
3.3 框架对比与选择：没有银弹
经常有人问我：“老师，PyTorch和TensorFlow到底选哪个？”我的回答永远是——看场景。
下面这张图是我自己总结的框架选择逻辑，画了很久才理清楚：
嗯，这张图其实反映了我多年的经验。简单总结一下：
| 维度 | PyTorch | TensorFlow |
|---|---|---|
| 调试体验 | ⭐ 动态图，print就能调试 | ⭐⭐ 静态图调试较麻烦 |
| 部署生态 | ⭐⭐ TorchServe + ONNX | ⭐⭐⭐ TF Serving + TFLite |
| 分布式训练 | ⭐⭐ DDP成熟，FSDP新 | ⭐⭐⭐ 原生支持完善 |
| 移动端支持 | ⭐ 相对较弱 | ⭐⭐⭐ TFLite很成熟 |
| 社区活跃度 | ⭐⭐⭐ 学术界首选 | ⭐⭐ 工业界为主 |
我的建议：
- 如果你做研究、写论文、快速验证想法——无脑选PyTorch。我带的实习生，第一天就能上手写模型。
- 如果你做工业级部署，尤其是移动端、IoT设备——TensorFlow更稳。它的TFLite和TF Serving经过大量生产验证。
- 如果你团队两种框架都熟——可以考虑混合使用：PyTorch做实验，转ONNX后用TensorRT推理。
注意：别盲目追新框架。我记得2020年JAX刚火的时候，很多人跟风迁移，结果发现文档不全、社区不成熟，最后又迁回PyTorch。选框架，稳定比时髦重要。
说到底，框架只是工具。我见过用PyTorch写出比TensorFlow还慢的代码，也见过用TensorFlow实现比PyTorch更灵活的模型。关键还是——你对你用的框架理解有多深。嗯，这一章的内容就到这儿，希望这些实战经验能帮你少走弯路。
想要更多资料，联系可可