深度学习框架实战对比：我在项目中的真实踩坑与收获
我最早接触深度学习的时候，TensorFlow 还是主流，PyTorch 算是小众新秀。那时候选框架更多是看社区和文档的丰富程度。但随着这几年在工业界项目的不断实践，我逐渐意识到：选择合适的深度学习框架，远不止“喜欢哪个语法风格”这么简单。
今天想借这篇文章，聊聊我在两个典型项目中分别使用 PyTorch 和 TensorFlow/Keras 的实际经验。不谈理论优劣，只讲实操中遇到的问题、踩过的坑以及最终的效果对比。希望通过我的亲身经历，能给正在选型或者对现有方案犹豫的朋友一点参考。
一、故事背景：两种场景下的技术选型挑战
项目1：图像分类任务（用的是Kaggle的一个花卉数据集）
我们团队接了一个客户定制的小项目，需求是训练一个花卉种类识别模型，用于他们App上的拍照识花功能。数据集是从Kaggle上找的公开花卉图像数据集（约34,000张图片，102类）。
当时我们需要快速验证几个骨干网络（ResNet、EfficientNet），同时考虑后续模型服务化上线的需求。这其实是一个标准的CV任务，适合使用高层封装好的工具来加速开发进度。
当时的疑问是：该用Keras还是PyTorch？我们团队成员大部分有PyTorch背景，但考虑到服务部署的流程更成熟，还是选择了Keras + TensorFlow。
项目2：文本生成任务（自建对话生成系统）
另一个项目是为某客服平台搭建一个基于Transformer的对话生成系统。目标是从用户历史对话中预测回复语句。数据来自真实业务积累的日志（经过脱敏处理）。这类任务属于NLP序列到序列建模，涉及较多自定义组件。
这里我们最终决定采用PyTorch。主要原因有两个：
- 需要高度灵活地控制模型结构（比如attention机制的细节）
- 我们计划引入一些前沿研究模型（如LoRA、Adapter等），需要更高的可定制性
这两个项目成了我对比两大主流框架的最佳实验场。
二、核心挑战：灵活性 vs 易用性的权衡
挑战1：模型迭代快，调试频繁（项目1）
我们在Kaggle花卉项目中遇到了几个典型的痛点：
- 初期尝试了多个预训练模型（通过
tf.keras.applications接口加载），非常方便； - 但在进行fine-tuning时，发现部分层无法按预期冻结（例如BatchNorm没有被冻结导致acc波动）；
- 自定义Loss函数时，TensorFlow的Eager Execution模式和Graph模式切换让人头疼。
举个例子，在实现Focal Loss时，由于TF 2.x默认的动态执行环境和静态图编译机制之间切换不当，出现了奇怪的梯度NaN问题。后来通过手动调用 .numpy() 强制转换解决了，但这显然不是优雅的方式。
挑战2：复杂逻辑控制困难（项目2）
在做文本生成模型的时候，我们尝试构建一个多阶段的训练流程：
- 先训Encoder部分
- 冻结Encoder再训Decoder
- 最后全量参数微调
这一过程中，PyTorch虽然给了我们自由去修改每个阶段的参数是否requires_grad，但也带来了额外的代码管理成本。特别是在梯度清零、累积、优化器分段初始化等操作上，稍不注意就出错。
记得有次因为把optimizer放在for循环外面，结果每次iteration optimizer都在复用前一次的参数组，模型根本不更新权重，愣是查了一天才发现问题所在 😭
三、具体实施方案：从设计思路到代码落地
项目1（Keras实现）关键代码片段
from tensorflow.keras.applications import ResNet50
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
base_model = ResNet50(weights='imagenet', include_top=False)
x = base_model.output
x = GlobalAveragePooling2D()(x)
predictions = Dense(102, activation='softmax')(x)
model = Model(inputs=base_model.input, outputs=predictions)
# 冻结前几层
for layer in base_model.layers[:100]:
layer.trainable = False
model.compile(optimizer=Adam(lr=1e-4), loss='categorical_crossentropy', metrics=['accuracy'])
model.fit(train_dataset, epochs=20, validation_data=val_dataset)
上面这段代码看起来挺干净简洁吧？但它背后隐藏了一些隐患。比如当使用混合精度训练时，如果你用了 tf.keras.mixed_precision.Policy("mixed_float16")，那么你可能还需要手动指定loss scale策略，否则训练容易爆炸。
另外，在做迁移学习时，一定要特别小心BN层的冻结方式，否则很容易让前面学到的特征在训练过程中被“洗掉”。
项目2（PyTorch实现）部分流程
from torch import nn
from transformers import BertModel
class Seq2Seq(nn.Module):
def __init__(self, encoder_name="bert-base-uncased", vocab_size=30522):
super().__init__()
self.encoder = BertModel.from_pretrained(encoder_name)
self.decoder = TransformerDecoder(...)
def forward(self, input_ids, attention_mask, decoder_input):
enc_out = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
dec_out = self.decoder(decoder_input, enc_out, ...)
model = Seq2Seq()
# 分阶段训练
optim_params = [
{"params": model.encoder.parameters(), "lr": 1e-5},
{"params": model.decoder.parameters(), "lr": 1e-3}
]
optimizer = torch.optim.AdamW(optim_params)
# 在train loop里根据条件toggle requires_grad
for param in model.encoder.parameters():
param.requires_grad = False # 第二阶段冻结encoder
在PyTorch中，我们可以非常细致地控制每一层的行为。尤其是在多任务学习或多阶段训练中，这种优势很明显。当然，这也意味着你要自己承担更多责任。
四、那些年我们一起踩过的坑
坑1：PyTorch的AMP自动混合精度失效
在项目2中，我们一开始尝试用torch.cuda.amp来做混合精度训练，以加快速度并节省显存。但训练过程中经常出现某些运算返回None，导致backward报错。
最后排查发现是因为有些自定义模块没写成支持AMP的形式。解决方法是检查所有forward中是否有非PyTorch原生的操作，并尽量使用 @torch.cuda.amp.autocast() 装饰器包裹相关函数。
坑2：Keras模型保存格式问题
Kaggle项目后期我们要导出ONNX进行推理优化。原本以为直接save_model就可以搞定，结果加载到ONNX Runtime时报错：“找不到某个layer的input_shape属性”。
原来是Keras默认保存的是SavedModel格式，而转换ONNX时需要明确输入shape。解决方案是在保存前先创建inputs层并调用build():
inputs = tf.keras.Input(shape=(224, 224, 3))
outputs = model(inputs)
final_model = tf.keras.Model(inputs=inputs, outputs=outputs)
final_model.save("flower_model.h5")
坑3：PyTorch的分布式训练性能瓶颈
我们在文本生成项目里用了多卡训练（DDP），但在启动脚本中设置device_ids时出了问题 —— 卡号顺序搞错了，导致每张卡都只读了自己的local_rank而不知道全局rank。
后来统一改用 torchrun --nproc_per_node=N 来管理进程，配合下面这行代码正确获取设备信息：
device = f'cuda:{dist.get_rank() % torch.cuda.device_count()}'
这个问题花了整整一天排查，教训就是：不要手动传rank参数，交给启动器来调度是最稳的。
五、效果对比：各有千秋，关键看场景
| 维度 | Keras/TensorFlow | PyTorch |
|---|---|---|
| 开发效率 | 快速搭建能力强，API统一 | 更加灵活，适合定制化开发 |
| 调试体验 | 图模型有时难以打断点调试 | Eager模式友好，变量可打印 |
| 上线部署 | TFServing集成好，转换ONNX/TF-Lite流畅 | ONNX支持不错，但需要额外工作 |
| 社区资源 | 文档丰富，教程多，案例齐全 | 学术研究首选，最新论文实现多 |
| 性能表现 | 多GPU训练更稳定 | 可控性高，但易出错 |
在这两个项目中，我们最终都达到了交付目标。Keras项目两周完成从数据预处理到模型部署，而PyTorch项目则用了一个月时间探索不同的架构设计和训练策略。
六、几点建议给同行朋友们
业务优先，技术其次。如果只是做一个标准任务（如图像分类、常规NLP分类），用Keras更省心；但如果要做科研性质的工作或创新模型，PyTorch更合适。
别盲目追求“高级感”。我见过不少朋友一味追求Transformer、Diffusion之类的大模型，却忽略了业务的核心指标要求。很多时候，轻量级的CNN也能达到95%+准确率，何必硬上BERT？
部署要考虑完整链路。训练模型只是一个环节，真正考验人的是如何把它变成产品可用的API。这时候你会发现Keras SavedModel + TF Serving真的香。
版本差异不容忽视。无论是TF 1.x到2.x，还是PyTorch 1.x到2.x，都有不小的breaking change。务必在初期锁定版本号。
保持开放心态。如今很多项目开始融合PyTorch和TF的优势，比如用PyTorch训练然后转成ONNX再部署到TF环境中运行。没必要死磕一种框架。
结语：深度学习不止框架之争
写到这里，我想起了几个月前的一次项目复盘会上我说过一句话：“框架本身不是目的，而是实现目标的手段。”这句话现在回头看，依旧适用。
无论你是钟情PyTorch的灵活性，还是习惯Keras的便捷高效，最终的目标都应该服务于业务需求和技术价值。在AI这条路上，没有“最牛”的工具，只有最适合的解决方案。
希望这篇来自一线开发者的真实分享，能在你面临技术选型时带来一些思考和启发。如果你也有类似的经验，欢迎留言一起探讨！
评论 0