TensorFlow 2.x 实战:从计算图、自动微分到模型训练与部署
1. 从零理解 TensorFlow它到底在算什么很多人第一次接触 TensorFlow脑子里冒出来的第一个问题不是“怎么用”而是“这东西到底在干嘛”。我刚开始也一样看到一堆张量、计算图、会话、占位符感觉像在看天书。后来真正跑通几个项目之后才明白TensorFlow 的核心其实就一句话把数学运算组织成一张可以自动求导、自动并行、自动部署的计算图然后高效地执行它。这句话听起来还是有点抽象我们拆开来看。假设你要算一个最简单的线性回归y wx b。在普通 Python 里你写y w * x bPython 立刻就算出了结果。但在 TensorFlow 里你写下的这行代码并不会马上执行而是先记录成一个“节点”等到你真正需要结果的时候才把整张图丢给底层引擎去跑。这个“先建图、后执行”的模式就是 TensorFlow 最底层的设计哲学。为什么要这么设计因为深度学习模型的运算量极大一次训练可能要跑几百万次矩阵乘法。如果每次都用 Python 解释器逐行执行速度会慢到无法忍受。TensorFlow 把计算图交给 C 底层引擎再调用 CPU、GPU 甚至 TPU 的并行能力才能把训练时间从“几天”压缩到“几小时”。这就是它和 NumPy 最本质的区别——NumPy 是即时计算TensorFlow 是延迟计算加自动微分。1.1 张量一切数据的统一容器TensorFlow 里所有的数据都叫“张量”Tensor。这个词听着高级其实就是一个多维数组。标量是 0 维张量向量是 1 维张量矩阵是 2 维张量彩色图片是 3 维张量高、宽、通道一批图片是 4 维张量批量大小、高、宽、通道。你可以把它理解成 NumPy 的ndarray但多了两个关键能力可以在 GPU 上运行以及可以自动记录梯度。我刚开始学的时候总想把张量当成普通数组来用结果踩了不少坑。比如张量的数据类型必须一致不能像 Python 列表那样混装整数和字符串张量的形状一旦确定大部分操作都要求形状匹配否则直接报错。这些限制一开始让人觉得麻烦但正是这些限制让 TensorFlow 能够在底层做大量优化。import tensorflow as tf # 创建一个常量张量 a tf.constant([[1, 2], [3, 4]], dtypetf.float32) b tf.constant([[5, 6], [7, 8]], dtypetf.float32) # 矩阵乘法 c tf.matmul(a, b) print(c) # 输出: tf.Tensor([[19. 22.] [43. 50.]], shape(2, 2), dtypefloat32)上面这段代码里tf.constant创建的是不可变张量。如果你需要参与训练的参数就得用tf.Variable。这两者的区别很关键常量在训练过程中不会变变量会被优化器不断更新。我见过不少新手把权重写成常量结果训练半天损失不下降排查半天才发现是这个问题。1.2 自动微分深度学习框架的灵魂如果说张量是 TensorFlow 的骨架那自动微分就是它的灵魂。深度学习的核心是反向传播而反向传播的本质就是链式法则求导。手动推导一个三层网络的梯度就已经让人头大更别说上百层的 ResNet 或者 Transformer。TensorFlow 的GradientTape就是干这个的你只需要写前向计算它自动帮你算所有参数的梯度。x tf.Variable(3.0) with tf.GradientTape() as tape: y x ** 2 2 * x 1 # 自动求导 dy_dx tape.gradient(y, x) print(dy_dx) # 输出: tf.Tensor(8.0, shape(), dtypefloat32)这个例子里y x² 2x 1 的导数是 2x 2代入 x3 得到 8。GradientTape会记录with块里所有涉及x的运算然后反向推导出梯度。这个机制看起来简单但它是所有神经网络训练的基础。我个人的经验是理解GradientTape之后再去看model.fit或者自定义训练循环就会觉得非常清晰因为你知道底层到底在发生什么。提示GradientTape默认只追踪一次梯度如果你需要求二阶导数要在with块里加上tape.watch(x)并设置persistentTrue。这个细节在实现物理信息神经网络或者对抗样本生成时会用到。1.3 计算图从静态到动态的演进TensorFlow 1.x 的时代计算图是静态的你先定义整张图再创建会话执行。这种方式效率高但调试极其痛苦因为报错信息往往指向图的结构而不是具体数据。TensorFlow 2.x 默认切换到了动态图模式Eager Execution写一行执行一行调试体验和普通 Python 几乎一样。但动态图并不意味着静态图消失了。当你用tf.function装饰一个函数时TensorFlow 会把它追踪成一张静态图从而获得更好的性能。这个机制叫 AutoGraph它会把 Python 的控制流if、for、while转换成 TensorFlow 的图操作。我实测下来对于一个包含复杂循环的自定义训练步骤加上tf.function之后训练速度能提升 30% 到 50%具体取决于模型结构和批量大小。tf.function def train_step(x, y): with tf.GradientTape() as tape: predictions model(x, trainingTrue) loss loss_fn(y, predictions) gradients tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) return loss这段代码第一次调用时会被追踪成图后续调用直接执行图省去了 Python 解释器的开销。但要注意tf.function对 Python 副作用比如打印、修改外部变量的支持有限调试时最好先用动态图跑通再切换到图模式加速。2. 环境搭建与第一个可训练模型搞清楚了核心概念接下来就是动手。TensorFlow 的环境搭建说简单也简单说坑多也坑多。我见过太多人卡在安装这一步要么版本不兼容要么 GPU 驱动对不上要么 import 之后一堆警告。这一章我把自己的安装流程和踩过的坑完整梳理一遍尽量让你一次跑通。2.1 安装方式的选择pip、conda 还是 Docker目前主流的安装方式有三种pip、conda 和 Docker。我的建议是如果你只是学习和做小项目直接用 pip 安装 CPU 版本就够了如果你有 NVIDIA 显卡并且需要训练较大模型conda 在管理 CUDA 和 cuDNN 依赖上更方便如果你需要部署到服务器或者保证环境完全可复现Docker 是最稳妥的选择。安装方式适用场景优点缺点pip学习、小项目简单直接版本最新GPU 依赖需手动配置conda科研、多环境自动处理 CUDA 依赖包版本可能滞后Docker部署、复现环境隔离彻底需要熟悉容器操作我个人的习惯是本地开发用 conda 创建独立环境服务器部署用 Docker。这样既能享受 conda 的便利又能保证线上环境的一致性。# 创建 conda 环境 conda create -n tf-env python3.10 conda activate tf-env # 安装 TensorFlowCPU 版本 pip install tensorflow # 验证安装 python -c import tensorflow as tf; print(tf.__version__)如果你需要 GPU 支持安装命令会稍微复杂一些。TensorFlow 2.x 之后GPU 版本和 CPU 版本合并了同一个包但你需要确保系统里安装了匹配版本的 CUDA 和 cuDNN。我踩过最大的坑就是 CUDA 版本和 TensorFlow 版本不匹配导致tf.config.list_physical_devices(GPU)返回空列表。解决办法是去 TensorFlow 官网查版本对应表严格按表格安装。注意安装完成后如果看到大量警告信息比如 “Could not load dynamic library”通常是 CUDA 路径没配好。如果只用 CPU可以在代码开头设置os.environ[TF_CPP_MIN_LOG_LEVEL] 3来屏蔽这些警告。2.2 用 Keras 快速搭建第一个神经网络TensorFlow 2.x 把 Keras 作为官方高阶 API这让搭建模型变得异常简单。我刚开始学的时候用原生 TensorFlow 写一个全连接网络要几十行代码现在用 Keras 只需要几行。但简单不代表可以跳过理解我建议你至少手动实现一次前向传播和反向传播再切换到 Keras这样遇到问题时才知道从哪里排查。import tensorflow as tf from tensorflow.keras import layers, models # 搭建一个简单的全连接网络 model models.Sequential([ layers.Dense(128, activationrelu, input_shape(784,)), layers.Dropout(0.2), layers.Dense(64, activationrelu), layers.Dense(10, activationsoftmax) ]) # 编译模型 model.compile( optimizeradam, losssparse_categorical_crossentropy, metrics[accuracy] ) # 查看模型结构 model.summary()这段代码搭建了一个三层全连接网络输入是 784 维对应 28x28 的灰度图像展平输出是 10 类概率。Dropout层用来防止过拟合softmax把输出转换成概率分布。model.summary()会打印每一层的参数数量你可以据此估算模型大小和显存占用。我特别想强调input_shape这个参数。很多新手会忽略它结果第一次调用model.fit时报错说形状不匹配。其实 Keras 可以在第一次前向传播时自动推断输入形状但显式指定能让模型结构更清晰也方便后续保存和加载。2.3 数据管道的构建tf.data 的正确用法模型搭好了接下来要喂数据。TensorFlow 提供了tf.dataAPI 来构建高效的数据管道。我见过很多人直接把 NumPy 数组丢给model.fit小数据集没问题但数据量一大就会成为瓶颈因为数据加载和模型计算是串行的。tf.data.Dataset的核心优势在于它可以预取prefetch、并行加载num_parallel_calls、缓存cache和批处理batch。这几个操作组合起来能让 GPU 利用率从 30% 提升到 90% 以上。# 假设 x_train 和 y_train 是 NumPy 数组 dataset tf.data.Dataset.from_tensor_slices((x_train, y_train)) dataset dataset.shuffle(buffer_size10000) dataset dataset.batch(32) dataset dataset.prefetch(tf.data.AUTOTUNE) # 训练时直接传入 dataset model.fit(dataset, epochs10)shuffle的buffer_size决定了打乱的程度太小会导致每个 epoch 的数据顺序相似太大会占用过多内存。我的经验是设置为数据集大小的 10% 到 20% 比较合适。prefetch(tf.data.AUTOTUNE)让 TensorFlow 自动决定预取多少个批次通常设置为AUTOTUNE就能获得不错的效果。还有一个容易被忽略的点如果数据集能放进内存加上.cache()可以避免每个 epoch 重新读取和预处理数据训练速度会有明显提升。但如果数据集太大缓存反而会导致内存溢出这时候可以考虑缓存到磁盘或者使用生成器。3. 训练过程中的关键参数与调优逻辑模型能跑起来只是第一步真正考验功力的是怎么让它跑得好。学习率设多少、批量大小选多大、优化器用哪个、过拟合怎么处理这些问题没有标准答案但有一些经过大量实践验证的经验法则。这一章我结合自己的项目经验把每个关键参数的调整逻辑讲清楚。3.1 学习率最重要的超参数没有之一如果只能调一个超参数那一定是学习率。学习率太大损失会震荡甚至发散学习率太小收敛慢到让人怀疑人生。我刚开始做项目时用默认的 0.001 跑所有模型结果有的模型三天都不收敛有的模型直接 NaN。后来才明白学习率必须根据模型深度、批量大小和优化器来调整。一个实用的策略是使用学习率预热warmup加余弦衰减cosine decay。预热是在训练初期把学习率从很小的值逐渐增加到目标值避免初期梯度不稳定余弦衰减是在训练后期逐渐降低学习率让模型更精细地收敛。initial_learning_rate 0.001 warmup_steps 1000 total_steps 50000 lr_schedule tf.keras.optimizers.schedules.CosineDecay( initial_learning_rateinitial_learning_rate, decay_stepstotal_steps, alpha0.01 ) optimizer tf.keras.optimizers.Adam(learning_ratelr_schedule)这段代码定义了一个余弦衰减的学习率计划最终学习率会降到初始值的 1%。alpha参数控制最低学习率设置为 0.01 意味着最低降到 0.00001。我实测下来这种策略在图像分类和自然语言处理任务上都很稳比固定学习率的效果好不少。提示如果你不确定学习率设多少合适可以用学习率范围测试LR Range Test。从极小的学习率开始每个批次指数增加记录损失变化损失下降最快的区间就是合适的学习率范围。这个方法我在多个项目里用过比盲目试错高效得多。3.2 批量大小与显存占用的平衡批量大小决定了每次更新梯度时用多少样本。批量越大梯度估计越准确但显存占用也越大批量越小梯度噪声越大但有时反而能帮助模型跳出局部最优。我的一般原则是在显存允许的范围内尽量选大然后用学习率来补偿。有一个经验公式批量大小翻倍学习率也翻倍。这是因为大批量意味着更准确的梯度方向可以用更大的步长。但这个关系不是线性的批量大到一定程度后继续增加带来的收益会递减。批量大小显存占用训练速度泛化能力适用场景16低慢较好小数据集、显存受限32中中好通用默认值64中高快一般中等数据集128高很快需调参大数据集、多卡训练256很高最快需特殊技巧大规模分布式训练我踩过的一个坑是在小数据集上用大批量结果模型泛化能力很差验证集准确率比训练集低十几个百分点。后来把批量从 256 降到 32同时把学习率从 0.01 降到 0.001泛化能力明显改善。所以批量大小不是越大越好要根据数据规模和任务难度来定。3.3 过拟合的识别与应对过拟合是训练中最常见的问题模型在训练集上表现很好在验证集上却很差。识别过拟合很简单看训练损失和验证损失的曲线。如果训练损失持续下降验证损失先下降后上升那就是过拟合了。应对过拟合的手段有很多我按优先级排序增加数据数据增强正则化Dropout早停。增加数据是最根本的解决办法但很多时候数据就那么多所以数据增强是性价比最高的选择。对于图像任务旋转、翻转、裁剪、颜色抖动都能显著提升泛化能力。# 数据增强示例 data_augmentation tf.keras.Sequential([ layers.RandomFlip(horizontal), layers.RandomRotation(0.1), layers.RandomZoom(0.1), ]) # 在模型中使用 model models.Sequential([ data_augmentation, layers.Conv2D(32, 3, activationrelu), layers.MaxPooling2D(), layers.Flatten(), layers.Dense(10, activationsoftmax) ])L2 正则化是另一种常用手段它通过在损失函数里加上权重的平方和来限制权重过大。tf.keras.regularizers.l2(0.001)可以直接加在层上。Dropout 则是在训练时随机丢弃一部分神经元强迫网络学习更鲁棒的特征。我通常会在全连接层前加 Dropout比例从 0.2 到 0.5 不等具体看模型大小。早停Early Stopping是最省事的办法监控验证损失如果连续几个 epoch 不下降就停止训练。Keras 提供了EarlyStopping回调设置patience5意味着验证损失连续 5 个 epoch 不改善就停止同时可以配合restore_best_weightsTrue恢复最佳权重。4. 模型保存、加载与部署的实战细节训练完模型只是完成了工作的一半怎么把它保存下来、加载到生产环境、高效地提供服务这些环节同样重要。我见过不少项目在实验室里跑得很好一到部署就各种问题模型文件太大、推理速度太慢、依赖版本冲突。这一章我把模型保存和部署的完整流程梳理一遍。4.1 SavedModel 与 H5 格式的选择TensorFlow 支持多种模型保存格式最常用的是 SavedModel 和 H5。SavedModel 是 TensorFlow 的原生格式保存了完整的计算图和权重跨平台兼容性最好适合部署到服务器、移动端或者浏览器。H5 格式是 Keras 的旧格式只保存模型结构和权重加载时需要相同的代码环境。# 保存为 SavedModel 格式 model.save(my_model) # 保存为 H5 格式 model.save(my_model.h5) # 加载模型 loaded_model tf.keras.models.load_model(my_model)我的建议是生产环境一律用 SavedModel。它不仅保存了权重还保存了前向传播的计算图加载后可以直接推理不需要重新定义模型结构。而且 SavedModel 支持签名signature可以定义输入输出的具体格式方便和其他系统集成。注意如果你在训练时用了自定义层或者自定义损失函数加载 SavedModel 时需要提供custom_objects参数否则会报错说找不到对应的类。这个坑我在第一次部署自定义模型时踩过排查了很久才发现是这个问题。4.2 推理优化从 float32 到量化模型训练时用 float32 精度是为了保证梯度计算的稳定性但推理时其实不需要这么高的精度。把模型量化成 int8 或者 float16可以显著减小模型体积、提升推理速度而精度损失通常很小。TensorFlow 提供了训练后量化Post-Training Quantization和量化感知训练Quantization-Aware Training两种方式。训练后量化最简单只需要一个代表数据集就能完成量化感知训练需要在训练时就模拟量化误差精度更高但流程更复杂。# 训练后量化示例 converter tf.lite.TFLiteConverter.from_saved_model(my_model) converter.optimizations [tf.lite.Optimize.DEFAULT] # 提供代表数据集 def representative_dataset(): for _ in range(100): yield [np.random.rand(1, 224, 224, 3).astype(np.float32)] converter.representative_dataset representative_dataset converter.target_spec.supported_ops [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] converter.inference_input_type tf.int8 converter.inference_output_type tf.int8 tflite_model converter.convert() with open(model_quantized.tflite, wb) as f: f.write(tflite_model)这段代码把 SavedModel 转换成 int8 量化的 TFLite 模型。我实测下来模型体积能缩小到原来的四分之一推理速度提升两到三倍精度损失通常在 1% 以内。对于移动端和嵌入式设备这个优化非常值得做。4.3 服务化部署TF Serving 与自定义 API如果要把模型部署成在线服务TensorFlow Serving 是官方推荐的方案。它支持模型热更新、版本管理、批量推理性能也很出色。但 TF Serving 的配置相对复杂需要理解模型签名、版本目录结构等概念。另一种更灵活的方式是用 Flask 或 FastAPI 自己写一个推理服务。这种方式适合小规模部署或者需要复杂预处理逻辑的场景。我通常会在服务启动时加载模型然后每个请求调用model.predict或者model.__call__。from fastapi import FastAPI import tensorflow as tf import numpy as np app FastAPI() model tf.keras.models.load_model(my_model) app.post(/predict) async def predict(data: dict): input_array np.array(data[input], dtypenp.float32) input_array np.expand_dims(input_array, axis0) prediction model.predict(input_array) return {prediction: prediction.tolist()}这个简单的 FastAPI 服务可以接收 JSON 格式的输入返回预测结果。实际部署时还需要考虑并发处理、请求队列、超时控制等问题。我的经验是如果 QPS 不高比如每秒几十个请求单进程加异步就够了如果 QPS 很高就需要用 TF Serving 或者多进程加负载均衡。提示模型加载是一个耗时操作千万不要在每个请求里加载模型。正确的做法是在服务启动时加载一次然后复用。我见过有人在请求处理函数里调用load_model结果每个请求都要等好几秒完全不可用。5. 那些文档里不会写的踩坑记录这一章我专门记录一些在实际项目中踩过的坑这些问题在官方文档里往往一笔带过但真正遇到时能让人折腾半天。我把排查过程和解决方案完整写出来希望能帮你节省一些时间。5.1 内存泄漏循环中创建模型导致的 OOM有一次我写了一个超参数搜索的脚本在循环里反复创建模型、训练、评估。跑了几十个组合之后程序突然崩溃报 OOM内存溢出。我一开始以为是数据加载的问题排查了半天才发现是 TensorFlow 的图没有释放。在 TensorFlow 2.x 里每次创建模型或者调用tf.function都会在底层创建计算图。如果在一个循环里反复创建这些图会累积在内存里最终导致溢出。解决办法是在每次循环结束时调用tf.keras.backend.clear_session()它会清理当前会话的图状态。for lr in [0.001, 0.01, 0.1]: tf.keras.backend.clear_session() model create_model() model.compile(optimizertf.keras.optimizers.Adam(lr)) model.fit(x_train, y_train, epochs5) # 评估...加上clear_session()之后内存占用稳定了循环跑几百次也不会崩溃。这个坑让我意识到TensorFlow 虽然用起来像普通 Python但底层还是有状态管理的不能完全当成无状态函数来用。5.2 GPU 显存碎片化为什么 batch size 调小也没用另一个让我头疼的问题是 GPU 显存碎片化。明明模型不大批量大小也调得很小但就是报 “Resource exhausted: OOM”。后来查资料才知道TensorFlow 默认会一次性申请所有 GPU 显存如果多个进程共享 GPU 或者反复创建销毁模型就会产生碎片。解决办法是开启显存按需增长gpus tf.config.experimental.list_physical_devices(GPU) if gpus: for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True)set_memory_growth让 TensorFlow 按需申请显存而不是一次性占满。这样多个进程可以共享 GPU碎片问题也缓解了。但要注意这个设置必须在创建任何张量之前调用否则不生效。还有一个相关的问题是如果你在 Jupyter Notebook 里反复运行同一个单元格旧的模型和计算图可能还占着显存。这时候需要重启内核才能彻底释放。我现在的习惯是在 Notebook 里做实验时每个实验结束后手动删除模型变量并调用clear_session()。5.3 数据管道的隐藏性能陷阱tf.data用起来简单但如果不注意细节性能可能比直接喂 NumPy 还差。我遇到过一个典型案例用Dataset.from_tensor_slices加载大量小文件训练速度极慢。排查后发现是数据加载成了瓶颈GPU 利用率只有 20%。问题的根源在于from_tensor_slices会把整个数据集加载到内存如果数据是文件路径列表每个批次都要重新读取文件。正确的做法是用Dataset.from_tensor_slices配合map函数并且设置num_parallel_calls来并行读取。def load_and_preprocess(path, label): image tf.io.read_file(path) image tf.image.decode_jpeg(image, channels3) image tf.image.resize(image, [224, 224]) image image / 255.0 return image, label dataset tf.data.Dataset.from_tensor_slices((image_paths, labels)) dataset dataset.map(load_and_preprocess, num_parallel_callstf.data.AUTOTUNE) dataset dataset.batch(32) dataset dataset.prefetch(tf.data.AUTOTUNE)num_parallel_callstf.data.AUTOTUNE让 TensorFlow 自动决定并行读取的线程数通常能充分利用多核 CPU。prefetch则让数据加载和模型计算重叠进行GPU 不用等数据。这两个设置加上之后我的训练速度提升了将近三倍。注意map函数里的操作最好是 TensorFlow 原生操作如果用 Python 函数会被 AutoGraph 转换成图操作但有些 Python 特性不支持。如果遇到奇怪的错误可以先关掉num_parallel_calls试试确认是并行化导致的问题。5.4 版本升级带来的 API 断裂TensorFlow 的版本迭代很快不同版本之间的 API 变化有时很大。我就遇到过从 2.4 升级到 2.10 之后原来能跑的代码报了一堆弃用警告有些甚至直接报错。比如tf.keras.optimizers.Adam的learning_rate参数在旧版本是lr新版本改成了learning_rate。我的建议是生产项目锁定版本在requirements.txt里写死 TensorFlow 的版本号不要用或者不指定版本。升级之前先在测试环境跑一遍完整流程确认所有 API 都兼容再上线。对于个人学习可以保持较新版本但要注意看官方迁移指南。还有一个常见问题是 CUDA 版本和 TensorFlow 版本的对应关系。TensorFlow 2.10 是最后一个支持 GPU 的原生 Windows 版本之后的版本在 Windows 上只能用 WSL2 或者 Docker。这个变化让很多 Windows 用户措手不及如果你在 Windows 上做深度学习建议尽早迁移到 WSL2 或者 Linux 环境。6. 从 TensorFlow 出发的进阶路线用了一段时间 TensorFlow 之后你可能会想接下来该学什么我的建议是不要急着学新框架先把 TensorFlow 的生态用透。TensorFlow 不只是一个训练框架它有一整套工具链TensorBoard 做可视化、TFX 做生产管道、TensorFlow Lite 做移动端部署、TensorFlow.js 做浏览器推理。这些工具组合起来能覆盖从实验到生产的完整流程。6.1 TensorBoard训练过程的可视化利器TensorBoard 是我用得最多的工具之一。它可以把训练损失、准确率、学习率、权重分布、计算图都可视化出来帮你快速判断模型是否在正常训练。用法很简单在训练时加一个回调就行。tensorboard_callback tf.keras.callbacks.TensorBoard( log_dir./logs, histogram_freq1, write_graphTrue ) model.fit(x_train, y_train, epochs10, callbacks[tensorboard_callback])然后在终端运行tensorboard --logdir ./logs打开浏览器就能看到各种图表。我特别推荐看权重分布图如果某一层的权重全部趋近于零或者爆炸式增长说明初始化或者学习率有问题。另外计算图可视化能帮你确认模型结构是否符合预期有时候层与层之间的连接错了光看代码不容易发现。6.2 自定义训练循环掌控每一个细节虽然model.fit很方便但有些场景需要更精细的控制比如对抗训练、多任务学习、自定义损失权重。这时候就需要写自定义训练循环。自定义循环的核心是GradientTape和optimizer.apply_gradients前面已经介绍过。tf.function def train_step(x_batch, y_batch): with tf.GradientTape() as tape: logits model(x_batch, trainingTrue) loss loss_fn(y_batch, logits) gradients tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) train_loss.update_state(loss) return loss for epoch in range(epochs): for x_batch, y_batch in train_dataset: loss train_step(x_batch, y_batch) print(fEpoch {epoch}, Loss: {train_loss.result()}) train_loss.reset_states()自定义循环的好处是你可以随时插入调试代码、修改梯度、调整损失权重。我做过一个多任务学习的项目两个任务的损失量级差很多用model.fit很难平衡改成自定义循环后给每个任务的损失加上可学习的权重效果好了很多。6.3 迁移学习站在预训练模型的肩膀上大多数实际项目不需要从零训练模型用预训练模型做迁移学习是更高效的选择。TensorFlow Hub 和 Keras Applications 提供了大量预训练模型涵盖图像分类、目标检测、文本分类、语义分割等任务。base_model tf.keras.applications.EfficientNetB0( include_topFalse, weightsimagenet, input_shape(224, 224, 3) ) base_model.trainable False model tf.keras.Sequential([ base_model, layers.GlobalAveragePooling2D(), layers.Dense(128, activationrelu), layers.Dropout(0.3), layers.Dense(num_classes, activationsoftmax) ])迁移学习的关键是先冻结预训练层只训练新加的分类头等分类头收敛后再解冻部分层做微调。微调时学习率要设得很小通常是初始学习率的十分之一甚至更低否则会破坏预训练学到的特征。我实测下来用 EfficientNet 做迁移学习在小数据集上几千张图片也能达到不错的准确率比从零训练快得多效果也好得多。提示选择预训练模型时不要只看准确率还要考虑推理速度和模型大小。EfficientNet 系列在准确率和效率之间平衡得很好MobileNet 适合移动端ResNet 是经典选择但参数量较大。根据你的部署环境来选不要盲目追求最高准确率。6.4 分布式训练从单卡到多卡当模型大到单卡放不下或者训练数据大到单卡跑不完时就需要分布式训练。TensorFlow 提供了MirroredStrategy做单机多卡MultiWorkerMirroredStrategy做多机多卡。用法很简单把模型创建和编译放在策略作用域里就行。strategy tf.distribute.MirroredStrategy() with strategy.scope(): model create_model() model.compile( optimizeradam, losssparse_categorical_crossentropy, metrics[accuracy] ) model.fit(train_dataset, epochs10)MirroredStrategy会自动把模型复制到每张 GPU同步梯度更新权重。批量大小要相应增大通常是单卡批量乘以 GPU 数量。我实测下来双卡训练速度大约是单卡的 1.8 倍四卡大约是 3.2 倍效率不是线性的因为梯度同步有通信开销。分布式训练最大的坑是数据管道。如果数据加载速度跟不上多卡也发挥不出优势。这时候需要确保tf.data管道有足够的并行度和预取必要时可以用strategy.experimental_distribute_dataset来分发数据。7. 写在最后一些个人体会TensorFlow 的学习曲线确实比较陡尤其是从 1.x 过渡到 2.x 的那段时间很多教程和代码都不兼容让人很困惑。但熬过最初的阶段之后你会发现它的设计其实很有逻辑张量是数据、计算图是执行、自动微分是核心、Keras 是接口、SavedModel 是交付。把这几个概念串起来整个框架就通了。我个人的经验是不要一上来就追求最新特性先把基础流程跑通加载数据、搭建模型、编译、训练、评估、保存、加载、推理。这个流程走顺了再去研究自定义层、自定义训练循环、分布式训练这些进阶内容。另外多动手写代码比看教程重要得多很多坑只有自己踩过才能记住。还有一个建议是保持对官方文档的关注。TensorFlow 的文档质量很高而且更新及时很多问题在文档里都有答案。遇到报错时先看错误信息的最后一行通常是关键然后去搜官方文档或者社区讨论大部分问题都能找到解决方案。最后不要被框架绑架。TensorFlow 和 PyTorch 都是工具核心的深度学习原理是相通的。把原理搞懂换框架只是换个 API 的事。我在实际项目中两个框架都用根据任务特点和团队习惯来选择没必要站队。重要的是解决问题而不是用什么工具。