有哪些相见恨晚的 TensorFlow 小技巧？

那些让我相见恨晚的 TensorFlow 小技巧

在数据科学和机器学习的世界里，TensorFlow 就像一位老朋友，陪伴我走过了不少征程。然而，随着项目越来越复杂，我总觉得在某个地方卡壳，优化效率、调试代码都耗费了不少精力。直到最近，我才挖掘出一些堪称“相见恨晚”的 TensorFlow 小技巧，它们就像隐藏的宝藏，一旦发现，便让我的开发体验焕然一新。今天，我就想好好跟大伙儿聊聊这些让我拍案叫绝的“秘籍”，希望也能帮到正在 TensorFlow 之路上摸索的你。

1. `tf.function`：别让 Python 成为性能瓶颈

坦白讲，最开始使用 TensorFlow 时，我其实挺依赖 Python 的动态性和灵活性，写起代码来顺畅无比。但是，当模型开始膨胀，训练数据量激增，我发现 Python 的解释执行模式成了一个巨大的性能瓶颈。神经网络的计算过程往往是海量的矩阵乘法和向量操作，这些操作交给 Python 来逐个执行，简直就是“慢如蜗牛”。

然后，我“偶然”发现了 `tf.function`。这个装饰器简直是救星！它能将 Python 函数“编译”成一个 TensorFlow 的计算图。一旦编译完成，后续的执行就如同在 C++ 或 Java 中一样高效。想象一下，一个原本需要几十秒才能跑完的训练步，在 `tf.function` 的加持下，可能只需要几秒钟。

怎么用？ 别以为它有多么复杂，只需要在你的模型训练函数、前向传播函数前加上 `@tf.function`。

```python
import tensorflow as tf
import time

假设这是你的训练步函数
@tf.function
def train_step(model, optimizer, inputs, targets):
with tf.GradientTape() as tape:
predictions = model(inputs, training=True)
loss = tf.keras.losses.categorical_crossentropy(targets, predictions)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return loss

模拟一个简单的模型
class SimpleModel(tf.keras.Model):
def __init__(self):
super(SimpleModel, self).__init__()
self.dense1 = tf.keras.layers.Dense(64, activation='relu')
self.dense2 = tf.keras.layers.Dense(10, activation='softmax')

def call(self, inputs):
x = self.dense1(inputs)
return self.dense2(x)

model = SimpleModel()
optimizer = tf.keras.optimizers.Adam()

第一次调用时会进行编译，之后就很快了
dummy_inputs = tf.random.normal((32, 100))
dummy_targets = tf.one_hot([i % 10 for i in range(32)], 10)

start_time = time.time()
for _ in range(100):
loss = train_step(model, optimizer, dummy_inputs, dummy_targets)
end_time = time.time()
print(f"Total time with tf.function: {end_time start_time:.4f} seconds")

对比一下没有 @tf.function 的情况（请自行尝试，会明显慢很多）
```

小提示：

首次调用慢： 第一次调用带有 `@tf.function` 的函数时，TensorFlow 会花一些时间进行追踪（tracing）和图的构建，所以会感觉比普通 Python 函数慢。但这是值得的，因为后续的调用会快很多。
`training=True/False`： 在模型的前向传播中，务必根据 `training` 参数来决定是否使用 `Dropout` 或 `BatchNormalization` 等会影响训练过程的层。`tf.function` 会为 `training=True` 和 `training=False` 分别生成不同的图，确保逻辑正确。
Python 变量陷阱： 在 `@tf.function` 内部，尽量避免使用 Python 的可变对象（如列表、字典）来存储模型训练过程中产生的中间结果，因为这些对象在图的追踪过程中可能导致问题。如果需要动态地收集信息，考虑使用 `tf.Tensor` 或者 `tf.Variable`。

2. `tf.data` API：数据流水线的艺术

处理大量数据时，我曾经遇到过数据加载速度跟不上模型训练速度的尴尬局面。这就像发动机再强劲，也得有足够的燃油供应。 `tf.data` API 就是我找到的“解决方案”。它提供了一套非常强大且灵活的工具，用于构建高效的数据输入流水线。

`tf.data` 的核心思想是将数据加载、预处理、批处理、乱序等操作解耦，并且可以利用多线程并行化，让数据准备的过程“润物细无声”。

关键组件：

`tf.data.Dataset.from_tensor_slices()`: 从 NumPy 数组或 TensorFlow 张量创建数据集。
`.map()`: 对数据集中的每个元素应用一个函数（例如，图像的缩放、归一化、数据增强）。
`.shuffle()`: 乱序数据集。
`.batch()`: 将数据分组成批次。
`.prefetch()`: 预取下一批数据，以便在当前批次训练时，下一批数据已经在 GPU/CPU 上准备好。这是提升吞吐量最关键的一步！
`.cache()`: 将预处理过的数据缓存在内存或磁盘中，避免重复计算。

示例：

```python
import tensorflow as tf
import os

假设我们有一些图像文件路径和对应的标签
image_paths = [f"path/to/image_{i}.jpg" for i in range(1000)]
labels = [i % 10 for i in range(1000)]

def load_and_preprocess_image(image_path, label):
这是一个占位符，实际应用中会读取图像文件
image = tf.io.read_file(image_path)
image = tf.image.decode_jpeg(image, channels=3)
image = tf.image.resize(image, [128, 128])
image = tf.image.convert_image_dtype(image, tf.float32)
简单的随机翻转作为数据增强
image = tf.image.random_flip_left_right(image)
return image, label

创建数据集
dataset = tf.data.Dataset.from_tensor_slices((image_paths, labels))

构建数据流水线
BUFFER_SIZE = 1000 乱序缓冲区大小
BATCH_SIZE = 32

dataset = dataset.map(load_and_preprocess_image, num_parallel_calls=tf.data.AUTOTUNE) 并行处理
dataset = dataset.cache() 缓存预处理结果
dataset = dataset.shuffle(BUFFER_SIZE)
dataset = dataset.batch(BATCH_SIZE)
dataset = dataset.prefetch(tf.data.AUTOTUNE) 预取数据

遍历数据集（在训练循环中使用）
for images, batch_labels in dataset:
训练模型...
pass
```

核心要义：

`tf.data.AUTOTUNE`： 这个值简直太好用了！它允许 TensorFlow 根据你的硬件情况自动调整并行度和预取数量，省去了手动调优的麻烦。
`prefetch()` 的魔力： `prefetch(tf.data.AUTOTUNE)` 是提升效率的关键。它让数据加载和模型训练可以并行进行，最大限度地利用 GPU/CPU 的计算能力。
`cache()` 的妙用： 如果你的数据集不大，或者预处理过程比较耗时，`cache()` 可以显著加快训练速度，避免重复执行预处理。
Chain Operations: `tf.data` 的 API 设计允许你通过链式调用来组合各种操作，非常直观。

3. `tf.keras.metrics`：告别手动统计

在模型训练过程中，跟踪准确率、损失值等指标是必不可少的。我之前总是手动在训练循环里写一堆代码来累加样本数、累加损失值，然后再计算平均值。这不仅繁琐，而且容易出错。

后来我才意识到，Keras 内置了强大的 `tf.keras.metrics` 模块，它们可以非常方便地统计和跟踪各种指标。

怎么用？

1. 实例化指标： 在训练开始前，创建你需要的指标实例。
2. 在训练循环中更新： 在每个训练步（或每个 epoch）结束后，使用 `metric.update_state(y_true, y_pred)` 来更新指标的状态。
3. 获取结果： 在需要的时候，使用 `metric.result()` 来获取当前的指标值。
4. 重置状态： 在每个 epoch 开始时，别忘了使用 `metric.reset_states()` 来清零指标，准备统计下一个 epoch 的数据。

示例：

```python
import tensorflow as tf

实例化指标
train_accuracy = tf.keras.metrics.CategoricalAccuracy()
train_loss = tf.keras.metrics.Mean()

假设这是你的训练循环
num_batches = 100
for i in range(num_batches):
模拟一批真实标签和模型预测
y_true = tf.one_hot([i % 5 for _ in range(32)], 5)
y_pred = tf.random.uniform((32, 5))

假设你计算了损失
loss = tf.reduce_mean(tf.square(y_true y_pred)) 随便找个损失函数

更新指标状态
train_accuracy.update_state(y_true, y_pred)
train_loss.update_state(loss)

（可选）在每个 batch 结束后打印结果
print(f"Batch {i+1}: Accuracy = {train_accuracy.result().numpy():.4f}, Loss = {train_loss.result().numpy():.4f}")

打印 epoch 结束时的最终结果
print(f"Epoch End: Final Accuracy = {train_accuracy.result().numpy():.4f}, Final Loss = {train_loss.result().numpy():.4f}")

重置指标，为下一个 epoch 做准备
train_accuracy.reset_states()
train_loss.reset_states()
```

重要提醒：

`reset_states()` 是关键！ 忘记这一步，你就会看到指标值不断累加，完全失去参考意义。
多样的指标： `tf.keras.metrics` 提供了各种各样的指标，比如 `Precision`, `Recall`, `F1Score`, `AUC` 等等，根据你的任务选择合适的指标。
自定义指标： 如果 Keras 没有你想要的指标，你也可以通过继承 `tf.keras.metrics.Metric` 类来创建自己的指标。

4. `tf.GradientTape` 的更高级用法：自定义训练循环的强大工具

虽然 Keras 的 `model.fit()` 方法很方便，但对于更复杂的训练逻辑，比如需要更精细地控制梯度，或者实现一些高级的训练技巧（如知识蒸馏、对抗训练），自定义训练循环就显得尤为重要。而 `tf.GradientTape` 就是实现这一目标的核心。

我之前总觉得 `GradientTape` 只能用来计算一个简单模型的梯度，但深入了解后，我才发现它的潜力和灵活性远超我的想象。

一些高级用法：

`watch()`： 默认情况下，`GradientTape` 会自动追踪所有在 `with` 块内计算的可训练变量。但如果你想追踪非变量的张量，或者某个特殊的变量，可以使用 `tape.watch()`。

```python
v = tf.Variable(0.0)
with tf.GradientTape() as tape:
y = v v
x = y + 1 x 不是变量，默认不会被追踪
tape.watch(x) 显式追踪 x
grad_x = tape.gradient(x, v) 这样也能计算出关于 v 的梯度
```

Persistent Tape： 默认情况下，`GradientTape` 在调用 `gradient()` 方法后就会被释放，只能计算一次梯度。如果你需要多次计算梯度（例如，计算二阶导数），需要将 `persistent=True` 传给 `GradientTape`。

```python
v = tf.Variable(0.0)
with tf.GradientTape(persistent=True) as tape:
y = v v
z = y y
grad_y = tape.gradient(y, v) 第一次计算
grad_z = tape.gradient(z, v) 第二次计算
del tape 当不需要时，释放 persistent tape
```

`tf.function` 结合 `GradientTape`： 如前所述，将自定义训练循环包装在 `@tf.function` 中，可以极大地提升训练速度。

```python
@tf.function
def custom_train_step(model, optimizer, inputs, targets):
with tf.GradientTape() as tape:
predictions = model(inputs, training=True)
loss = tf.keras.losses.MeanSquaredError()(targets, predictions)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return loss

... (模型、优化器、数据准备与前面类似)
在训练循环中调用 custom_train_step
```

我的感悟：

`tf.GradientTape` 就像一个“全能工具箱”，你可以在里面自由地计算各种梯度，从而实现各种复杂的训练逻辑。而 `@tf.function` 则为这个工具箱安装了“涡轮增压器”，让你的自定义训练循环效率飞起。

5. TensorBoard：让你的模型“开口说话”

模型训练过程中，我们经常需要可视化损失曲线、准确率变化、计算图结构，甚至模型权重分布。手动去画图不仅耗时，而且不够直观。

TensorBoard 就是 TensorFlow 的“内置可视化大杀器”。它能够收集训练过程中产生的各种日志信息，并将其以图表、仪表盘等多种形式展示出来，让你能够清晰地了解模型的训练状况。

如何使用？

1. 设置 `tf.keras.callbacks.TensorBoard`： 在 `model.fit()` 中，你可以直接传入 `TensorBoard` 回调。

```python
import tensorflow as tf
from tensorflow.keras.callbacks import TensorBoard

... (模型、数据准备)

log_dir = "logs/fit/" + datetime.datetime.now().strftime("%Y%m%d%H%M%S")
tensorboard_callback = TensorBoard(log_dir=log_dir, histogram_freq=1) histogram_freq=1 表示每 epoch 记录一次权重直方图

model.fit(x_train, y_train, epochs=10, validation_split=0.2,
callbacks=[tensorboard_callback])
```

2. 手动写入日志： 如果是自定义训练循环，你可以使用 `tf.summary` 模块手动记录日志。

```python
import datetime
writer = tf.summary.create_file_writer("logs/custom_training")

在训练循环中
with writer.as_default():
tf.summary.scalar('loss', loss_value, step=epoch)
tf.summary.scalar('accuracy', accuracy_value, step=epoch)
记录直方图（例如，某个层的权重）
tf.summary.histogram('layer_weights', model.layers[0].get_weights()[0], step=epoch)
writer.flush()
```

启动 TensorBoard：

在你的终端中，导航到你的项目根目录，然后运行：

```bash
tensorboard logdir logs/fit
```

然后在浏览器中访问 TensorBoard 提供的地址（通常是 `http://localhost:6006/`）。

为什么我后来才重视它？

一开始，我总觉得“能跑通就行”，可视化是锦上添花。但当我遇到模型不收敛、训练速度异常慢等问题时，才发现没有 TensorBoard， debugging 简直是“大海捞针”。它能让你看到模型在“做什么”，而不是仅仅知道它“做了什么”。

结语

这些小技巧，虽然有些听起来可能很简单，但它们所带来的效率提升和代码优化是实实在在的。从将 Python 瓶颈转移到 TensorFlow 的计算图，到构建高效的数据流水线，再到利用 TensorBoard 来“读懂”模型，每一个都让我觉得“早点知道就好了”。

如果你也还在 TensorFlow 的旅程中，希望这些分享能为你带来一些启发。记住，学习和实践是不断进步的关键，而那些“相见恨晚”的技巧，往往就藏在你每一次深入的探索之中。

介绍两个很有用的技巧

1. 使用timeline来优化优化性能，timeline可以分析整个模型在forward和backward的时候,每个操作消耗的时间，由此可以针对性的优化耗时的操作。我之前尝试使用tensorflow多卡来加速训练的时候， 最后发现多卡速度还不如单卡快，改用tf.data来 加速读图片还是很慢，最后使用timeline分析出了速度慢的原因，timeline的使用如下

       run_metadata = tf.RunMetadata() run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE) config = tf.ConfigProto(graph_options=tf.GraphOptions(         optimizer_options=tf.OptimizerOptions(opt_level=tf.OptimizerOptions.L0))) with tf.Session(config=config) as sess:     c_np = sess.run(c,options=run_options,run_metadata=run_metadata)     tl = timeline.Timeline(run_metadata.step_stats)     ctf = tl.generate_chrome_trace_format() with open('timeline.json','w') as wd:     wd.write(ctf)     

然后到谷歌浏览器中打卡chrome://tracing 并导入 timeline.json ，最后可以看得如下图所示的每个操作消耗的时间，

这里横坐标为时间，从左到右依次为模型一次完整的forward and backward过程中，每个操作分别在cpu,gpu 0, gpu 1上消耗的时间，这些操作可以放大，非常方便观察具体每个操作在哪一个设备上消耗多少时间。这里我们cpu上主要有QueueDequeue操作，这是进行图片预期过程，这个时候gpu在并行计算的所以gpu没有空等；另外我的模型还有一个PyFunc在cpu上运行，如红框所示，此时gpu在等这个结果，没有任何操作运行，这个操作应该要优化的。另外就是如黑框所示，gpu上执行的时候有很大空隙，如黑框所示，这个导致gpu上的性能没有很好的利用起来，最后分析发现是我bn在多卡环境下没有使用正确，bn有一个参数updates_collections我设置为None 这时bn的参数mean,var是立即更新的，也是计算完当前layer的mean,var就更新，然后进行下一个layer的操作，这在单卡下没有问题的， 但是多卡情况下就会写等读的冲突，因为可能存在gpu0更新（写）mean但此时gpu1还没有计算到该层，所以gpu0就要等gpu1读完mean才能写，这样导致了 如黑框所示的空隙，这时只需将参数设置成updates_collections=tf.GraphKeys.UPDATE_OPS 即可，表示所以的bn参数更新由用户来显示指定更新，如

       update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)   with tf.control_dependencies(update_ops):     train_op = optimizer.minimize(loss)     

这样可以在每个卡forward完后，再更新bn参数，此时的写写不存在冲突。优化后，我的2卡训练获得了接近2倍的加速比。

2. 使用check= tf.addchecknumerics_ops,sess.run([check, ...])来检查NaN问题 ,该操作会报告所有出现NaN的操作，从而方便找到NaN的源头。