import tensorflow as tf
import time
import os

# 确保结果可复现（可选）
tf.keras.utils.set_random_seed(42)

# --- 1. 加载并准备数据 ---
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()

# 将像素值归一化到 [0, 1] 并转换为 float32 类型
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0

# 将标签转换为独热编码
y_train = tf.keras.utils.to_categorical(y_train, 10)
y_test = tf.keras.utils.to_categorical(y_test, 10)

# 创建一个基准 tf.data 管道
BATCH_SIZE = 128
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=1024).batch(BATCH_SIZE)

# --- 2. 定义一个简单的 CNN 模型 ---
def build_model():
    model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=(32, 32, 3)),
        tf.keras.layers.Conv2D(32, (3, 3), activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dense(10) # 输出 logits (线性激活)
    ])
    return model

# --- 3. 基准训练步骤 ---
# 使用带 from_logits=True 的交叉熵损失
loss_fn = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.Adam()

# 我们将使用 tf.function 作为基本的图优化基准
@tf.function
def train_step(images, labels, model):
    with tf.GradientTape() as tape:
        predictions = model(images, training=True)
        loss = loss_fn(labels, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

# --- 4. 基准训练循环 ---
def run_training(dataset, model, steps_per_epoch, epochs, profile=False, logdir=None):
    if profile and logdir:
        tf.profiler.experimental.start(logdir)

    total_steps = steps_per_epoch * epochs
    step_count = 0
    start_time = time.time()

    for epoch in range(epochs):
        print(f"Epoch {epoch+1}/{epochs}")
        epoch_start_time = time.time()
        epoch_loss_avg = tf.keras.metrics.Mean()

        for step, (images, labels) in enumerate(dataset):
            if step_count == 10 and profile and logdir: # 在热身阶段后分析几个步骤
                 tf.profiler.experimental.stop()
                 print(f"Profiler data saved to {logdir}")
                 profile=False # 避免再次停止

            loss = train_step(images, labels, model)
            epoch_loss_avg.update_state(loss)
            step_count += 1

            if step >= steps_per_epoch:
                break

        epoch_time = time.time() - epoch_start_time
        steps_time = epoch_time / steps_per_epoch
        print(f"  Steps/sec: {1.0/steps_time:.2f}, Avg Loss: {epoch_loss_avg.result():.4f}, Time: {epoch_time:.2f}s")

    total_time = time.time() - start_time
    print(f"\nTotal Training Time: {total_time:.2f}s")
    avg_step_time = total_time / total_steps if total_steps > 0 else 0
    return avg_step_time

# --- 运行基准训练 ---
print("Running Baseline Training...")
baseline_model = build_model()
STEPS_PER_EPOCH = len(x_train) // BATCH_SIZE // 4 # 为加快测试使用较少的步骤
EPOCHS = 3

baseline_step_time = run_training(train_dataset, baseline_model, STEPS_PER_EPOCH, EPOCHS)
print(f"\nBaseline Average Step Time: {baseline_step_time*1000:.2f} ms")

在继续之前，请执行此代码。注意最后打印出的平均步长耗时。这是我们的基准性能指标。

使用 TensorBoard 进行性能分析

现在，我们使用 TensorBoard Profiler 来了解时间花在哪里。我们需要稍微修改训练循环，以便对部分步骤进行性能分析。

创建一个日志目录： 这是性能分析数据存储的位置。
使用 tf.profiler.experimental.start/stop 包裹训练步骤： 我们将在初始热身阶段后分析几个步骤。

修改 run_training 调用：

# --- 运行启用性能分析的训练 ---
print("\nRunning Training with Profiling...")
profile_model = build_model() # 使用一个新的模型实例
LOG_DIR = "./logs/profile_baseline"
os.makedirs(LOG_DIR, exist_ok=True)

# 重新运行，启用 profile=True 并指定 logdir
run_training(train_dataset, profile_model, STEPS_PER_EPOCH, EPOCHS, profile=True, logdir=LOG_DIR)

运行此代码后，启动 TensorBoard：

tensorboard --logdir ./logs

在浏览器中导航到“Profile”选项卡（通常在 http://localhost:6006/）。查看这些工具：

概览页面： 提供高级别概述。检查“输入管道分析”，它可能显示您的 CPU 是否难以跟上 GPU 的速度（输入延迟高）。
追踪查看器： 显示 CPU 和 GPU 上操作的详细时间线。查找 GPU 空闲的空白区域，这可能表示它在等待数据。
TensorFlow 统计： 列出各个 TensorFlow 操作的执行时间。有助于发现意外慢的操作。

解读： 假设性能分析器显示大量时间花费在输入处理上（“输入管道分析”显示高延迟），并且 GPU 追踪查看器显示 GPU 在步骤之间有空闲期。这表明我们的 tf.data 管道是性能瓶颈。

优化 1：改进输入管道

根据我们的性能分析结果，让我们使用 .cache() 和 .prefetch() 优化 tf.data 管道。.cache() 在第一个 epoch 后将初始数据集元素保留在内存中，而 .prefetch() 则使数据预处理和模型执行重叠。

# --- 优化后的 tf.data 管道 ---
def create_optimized_dataset(x, y, batch_size):
    dataset = tf.data.Dataset.from_tensor_slices((x, y))
    dataset = dataset.shuffle(buffer_size=1024)
    dataset = dataset.batch(batch_size)
    # 添加缓存和预取
    dataset = dataset.cache() # 如果数据集适合内存，则在批处理后进行缓存
    dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
    return dataset

print("\nRunning Training with Optimized Dataset...")
optimized_dataset = create_optimized_dataset(x_train, y_train, BATCH_SIZE)
opt_data_model = build_model() # 新模型

opt_data_step_time = run_training(optimized_dataset, opt_data_model, STEPS_PER_EPOCH, EPOCHS)
print(f"\nOptimized Dataset Average Step Time: {opt_data_step_time*1000:.2f} ms")
print(f"Improvement vs Baseline: {(baseline_step_time / opt_data_step_time):.2f}x")

运行此更新的代码。您应该会观察到平均步长耗时显著减少，证实输入管道确实是限制因素。具体改进程度取决于您的硬件（CPU 速度、内存带宽）。

优化 2：启用混合精度训练

如果您有兼容的 GPU（NVIDIA Volta、Turing、Ampere 架构或更新的），混合精度可以通过在可能的情况下使用 16 位浮点数 ( $float16$ ) 进行计算，同时使用 $float32$ 保持模型精度，用于某些关键部分（例如变量更新），从而提供显著的加速和内存节省。

import tensorflow as tf
# 全局启用混合精度
# 在构建模型*之前*执行此操作
tf.keras.mixed_precision.set_global_policy('mixed_float16')

print("\nRunning Training with Mixed Precision...")

# --- 设置策略*后*重建模型 ---
# Keras 层将自动适应全局策略
mixed_precision_model = build_model()

# 重新编译优化器（如果需要损失缩放，Keras 会自动进行包装）
# 通常情况下，使用相同的损失和优化器实例是可以的，
# 但如果需要，重新创建它们可以确保状态干净。
optimizer_mp = tf.keras.optimizers.Adam()
loss_fn_mp = tf.keras.losses.CategoricalCrossentropy(from_logits=True)

# 修改后的设置需要一个新的 tf.function
@tf.function
def train_step_mp(images, labels, model, loss_fn, optimizer):
    with tf.GradientTape() as tape:
        predictions = model(images, training=True)
        # 如果需要，确保预测结果为 float32 以进行损失计算
        # Keras 层通常会处理此问题，但损失函数默认期望 float32。
        loss = loss_fn(labels, tf.cast(predictions, tf.float32))

    # 当使用混合精度策略时，损失缩放由 Keras 自动处理
    # 如果优化器是在设置策略后创建或重新编译的。
    # 对于自定义循环，您可能需要 `optimizer.get_scaled_loss(loss)` 和
    # `optimizer.get_unscaled_gradients(gradients)`。
    # 然而，使用标准 Keras model.fit 或这种带有
    # `apply_gradients` 的 @tf.function 方法，Keras 会处理。
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

# --- 混合精度的修改训练循环 ---
# (本质上是相同的循环，但确保使用混合精度模型、损失和优化器)
def run_training_mp(dataset, model, loss_fn, optimizer, steps_per_epoch, epochs):
    total_steps = steps_per_epoch * epochs
    step_count = 0
    start_time = time.time()

    for epoch in range(epochs):
        print(f"Epoch {epoch+1}/{epochs}")
        epoch_start_time = time.time()
        epoch_loss_avg = tf.keras.metrics.Mean()

        for step, (images, labels) in enumerate(dataset):
            loss = train_step_mp(images, labels, model, loss_fn, optimizer) # 使用混合精度训练步骤
            epoch_loss_avg.update_state(loss)
            step_count += 1
            if step >= steps_per_epoch:
                break

        epoch_time = time.time() - epoch_start_time
        steps_time = epoch_time / steps_per_epoch
        print(f"  Steps/sec: {1.0/steps_time:.2f}, Avg Loss: {epoch_loss_avg.result():.4f}, Time: {epoch_time:.2f}s")

    total_time = time.time() - start_time
    print(f"\nTotal Training Time: {total_time:.2f}s")
    avg_step_time = total_time / total_steps if total_steps > 0 else 0
    return avg_step_time

# 使用之前优化过的数据集
mixed_precision_step_time = run_training_mp(optimized_dataset, mixed_precision_model, loss_fn_mp, optimizer_mp, STEPS_PER_EPOCH, EPOCHS)

print(f"\nMixed Precision Average Step Time: {mixed_precision_step_time*1000:.2f} ms")
print(f"Improvement vs Optimized Data: {(opt_data_step_time / mixed_precision_step_time):.2f}x")
print(f"Improvement vs Baseline: {(baseline_step_time / mixed_precision_step_time):.2f}x")

# 如果之后运行其他需要 float32 的代码，则重置策略
# tf.keras.mixed_precision.set_global_policy('float32')

执行此代码。如果您的硬件支持高效的 $float16$ 计算（使用 NVIDIA GPU 上的 Tensor Cores），您应该会看到又一次性能提升。注意，第一个 epoch 可能会因为初始开销而稍微慢一些，但随后的 epoch 应该会更快。

优化 3：启用 XLA 编译

XLA（加速线性代数）是针对线性代数的领域特定编译器，它可以将多个 TensorFlow 操作合并为更高效、特定于硬件的内核。它可以在 GPU 和 CPU 上提供加速，尽管其效果通常在 TPU 和 GPU 上最显著。我们可以通过在 tf.function 装饰器中添加 jit_compile=True 来启用它。

# --- 定义 XLA 训练步骤 ---
@tf.function(jit_compile=True) # 启用 XLA
def train_step_xla(images, labels, model, loss_fn, optimizer):
    # 注意：这里也使用了混合精度策略
    with tf.GradientTape() as tape:
        predictions = model(images, training=True)
        loss = loss_fn(labels, tf.cast(predictions, tf.float32))
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

# --- XLA 的修改训练循环 ---
# (结构相似，只是使用 XLA 编译的训练步骤)
def run_training_xla(dataset, model, loss_fn, optimizer, steps_per_epoch, epochs):
    total_steps = steps_per_epoch * epochs
    step_count = 0
    start_time = time.time()

    # 执行一个初始步骤以在计时之外触发 XLA 编译
    print("Compiling XLA function (first step may be slow)...")
    _ = train_step_xla(next(iter(dataset))[0], next(iter(dataset))[1], model, loss_fn, optimizer)
    print("Compilation finished.")

    for epoch in range(epochs):
        print(f"Epoch {epoch+1}/{epochs}")
        epoch_start_time = time.time()
        epoch_loss_avg = tf.keras.metrics.Mean()

        for step, (images, labels) in enumerate(dataset):
            loss = train_step_xla(images, labels, model, loss_fn, optimizer) # 使用 XLA 步骤
            epoch_loss_avg.update_state(loss)
            step_count += 1
            if step >= steps_per_epoch:
                break

        epoch_time = time.time() - epoch_start_time
        steps_time = epoch_time / steps_per_epoch
        print(f"  Steps/sec: {1.0/steps_time:.2f}, Avg Loss: {epoch_loss_avg.result():.4f}, Time: {epoch_time:.2f}s")

    total_time = time.time() - start_time - (time.time() - epoch_start_time) # 调整计时循环外花费的时间
    print(f"\nTotal Training Time (post-compile): {total_time:.2f}s")
    avg_step_time = total_time / total_steps if total_steps > 0 else 0
    return avg_step_time

print("\nRunning Training with Mixed Precision + XLA...")
# 继续使用混合精度策略和优化后的数据集
# 如果层状态受到先前运行的影响，我们需要一个新的模型实例，
# 或者如果合适，直接继续使用 mixed_precision_model。
# 为简单起见，这里我们假设继续使用 mixed_precision_model。
# 如果出现问题，重建模型：xla_model = build_model()
xla_model = mixed_precision_model # 重用使用混合精度训练的模型
optimizer_xla = tf.keras.optimizers.Adam() # 可能会重置优化器状态
loss_fn_xla = tf.keras.losses.CategoricalCrossentropy(from_logits=True)

xla_step_time = run_training_xla(optimized_dataset, xla_model, loss_fn_xla, optimizer_xla, STEPS_PER_EPOCH, EPOCHS)
print(f"\nMixed Precision + XLA Average Step Time: {xla_step_time*1000:.2f} ms")
print(f"Improvement vs Mixed Precision Only: {(mixed_precision_step_time / xla_step_time):.2f}x")
print(f"Improvement vs Baseline: {(baseline_step_time / xla_step_time):.2f}x")

# 如果完成，重置策略
tf.keras.mixed_precision.set_global_policy('float32')

运行此代码时，您会注意到第一个 epoch 真正开始之前可能存在延迟。这是 XLA 编译时间。如果 XLA 成功优化了计算图，后续步骤应该会执行得更快。XLA 的效果很大程度上取决于模型结构和硬件。有时，对于简单模型，开销可能超过收益，而对于包含许多可合并操作的复杂模型，收益可能会非常显著。

结果概览

让我们直观地看看改进。我们收集了每个阶段的平均步长耗时：

基准
优化后的 tf.data 管道
混合精度（基于优化数据）
XLA（基于混合精度和优化数据）

不同优化阶段的平均训练步长耗时（毫秒）。越低越好。（注意：实际值取决于您的硬件和执行期间观察到的具体时间。）

此实践练习体现了一种系统的性能调优方法：

设定基准： 在优化前测量性能。
性能分析： 使用 TensorBoard Profiler 等工具找出性能瓶颈。
应用针对性优化： 解决找出的瓶颈（例如，输入管道、计算）。优化后的 tf.data、混合精度和 XLA 等技术是功能强大的工具。
测量与比较： 量化每项优化的效果。

请记住，性能优化通常是迭代的。应用一项优化后，瓶颈可能会转移，需要进一步的性能分析和调整。最佳的技术组合取决于您的具体模型、数据集和硬件配置。尝试是获得最大性能的必要步骤。

这部分内容有帮助吗？