启用Float16混合精度后训练变慢问题求助（TensorFlow 2.10）

在Windows系统中使用TensorFlow 2.10搭配NVIDIA RTX 2060 SUPER（带张量核心）训练CIFAR-100时，启用Float16混合精度后每轮训练耗时反而比常规float32更长，但在TensorFlow 2.15版本的Google Colab中运行相同代码，混合精度能正常加速，符合预期。

复现代码 #

import tensorflow as tf
import ssl

ssl._create_default_https_context = ssl._create_unverified_context

(train_x, train_y), (test_x, test_y) = tf.keras.datasets.cifar100.load_data()

tf.keras.mixed_precision.set_global_policy("mixed_float16")

model = tf.keras.Sequential([
    
    tf.keras.layers.Lambda(lambda x : x / 255, input_shape=(32,32,3)),
    tf.keras.layers.Conv2D(filters=64, kernel_size=(4,4)),
    tf.keras.layers.MaxPool2D(),
    tf.keras.layers.Conv2D(filters=32, kernel_size=(2,2)),
    tf.keras.layers.MaxPool2D(),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(4096, activation="relu"),
    tf.keras.layers.Dense(4096, activation="relu"),
    tf.keras.layers.Dense(4096, activation="relu"),
    tf.keras.layers.Dense(4096, activation="relu"),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(100),
    tf.keras.layers.Activation("softmax", dtype="float32")
    ])

model.compile(optimizer="adam", loss=tf.keras.losses.SparseCategoricalCrossentropy(), metrics=["accuracy"])

print("compute dtype of first layer: ", model.layers[0].compute_dtype)

model.fit(train_x, train_y, epochs=100, batch_size=1020)

model.evaluate(test_x, test_y)

环境与现象 #

• 本地环境：Windows + TensorFlow 2.10 + CUDA 11.2 + cuDNN 8.1.1
• 现象对比：
  • 未启用混合精度：每轮训练耗时较短
  • 启用混合精度：每轮训练耗时显著增加

排查与解决建议 #

1. 手动启用动态损失缩放
   TF 2.10的自动损失缩放在Windows环境下可能存在适配问题，尝试手动配置动态损失缩放：

   policy = tf.keras.mixed_precision.Policy('mixed_float16', loss_scale='dynamic')
   tf.keras.mixed_precision.set_global_policy(policy)
   

2. 验证张量核心使用率
   使用NVIDIA Nsight Systems或nvidia-smi dmon工具监控训练过程中的张量核心使用率。若使用率为0，说明运算未被调度到张量核心，可能是TF 2.10对RTX 20系列的张量核心优化不足。

3. 调整batch size与模型结构
   当前模型包含多个4096神经元的Dense层，batch size=1020可能无法充分利用张量核心的并行能力。尝试增大batch size（如2048，需保证显存足够），或减少Dense层的神经元数量，观察性能变化。

4. 升级TensorFlow版本
   TF 2.10到2.15之间修复了大量混合精度相关的bug，尤其是Windows环境下的CUDA适配问题：

   • 若保留CUDA 11.2，可升级到TF 2.12（官方兼容版本）
   • 若允许升级CUDA，可直接升级到TF 2.15 + CUDA 11.8

5. 检查CUDA环境配置
   确认系统环境变量CUDA_PATH、CUDA_PATH_V11_2正确设置，且cuDNN的bin、include、lib目录已添加到系统PATH中。运行以下代码验证GPU加速状态：

   print(tf.config.list_physical_devices('GPU'))
   print(tf.test.is_built_with_cuda())
   print(tf.test.is_gpu_available(cuda_only=True))
   

内容的提问来源于stack exchange，提问作者Space Programmer