TensorFlow数据增强出现while_loop转换警告的原因及解决方法

我按照TensorFlow官方教程实现了数据增强，代码功能正常，但运行时出现大量Using a while_loop for converting X的警告，使用的TensorFlow版本为v2.9.1。

实现代码 #

数据增强层定义 #

def _getAugmentationFunction(self):
    if not self.augmentation:
        return None
    pipeline = []
    
    pipeline.append(layers.RandomFlip('horizontal_and_vertical'))
    pipeline.append(layers.RandomRotation(30))
    pipeline.append(layers.RandomTranslation(0.1, 0.1, fill_mode='nearest'))
    pipeline.append(layers.RandomBrightness(0.1, value_range=(0.0, 1.0)))

    model =  Sequential(pipeline)
    return lambda x, y: (model(x, training=True), y)

应用数据增强到数据集 #

data_augmentation = self._getAugmentationFunction()
self.train_data = self.train_data.map(data_augmentation,
                                      num_parallel_calls=AUTOTUNE)

警告信息 #

WARNING:tensorflow:Using a while_loop for converting RngReadAndSkip
WARNING:tensorflow:Using a while_loop for converting Bitcast
WARNING:tensorflow:Using a while_loop for converting Bitcast
WARNING:tensorflow:Using a while_loop for converting StatelessRandomUniformV2
WARNING:tensorflow:Using a while_loop for converting RngReadAndSkip
WARNING:tensorflow:Using a while_loop for converting Bitcast
WARNING:tensorflow:Using a while_loop for converting Bitcast
WARNING:tensorflow:Using a while_loop for converting StatelessRandomUniformV2
WARNING:tensorflow:Using a while_loop for converting ImageProjectiveTransformV3
WARNING:tensorflow:Using a while_loop for converting RngReadAndSkip
WARNING:tensorflow:Using a while_loop for converting Bitcast
WARNING:tensorflow:Using a while_loop for converting Bitcast
WARNING:tensorflow:Using a while_loop for converting StatelessRandomUniformV2
WARNING:tensorflow:Using a while_loop for converting RngReadAndSkip
WARNING:tensorflow:Using a while_loop for converting Bitcast
WARNING:tensorflow:Using a while_loop for converting Bitcast
WARNING:tensorflow:Using a while_loop for converting StatelessRandomUniformV2
WARNING:tensorflow:Using a while_loop for converting RngReadAndSkip
WARNING:tensorflow:Using a while_loop for converting Bitcast
WARNING:tensorflow:Using a while_loop for converting Bitcast
WARNING:tensorflow:Using a while_loop for converting StatelessRandomUniformV2
WARNING:tensorflow:Using a while_loop for converting ImageProjectiveTransformV3
WARNING:tensorflow:Using a while_loop for converting RngReadAndSkip
WARNING:tensorflow:Using a while_loop for converting Bitcast
WARNING:tensorflow:Using a while_loop for converting Bitcast
WARNING:tensorflow:Using a while_loop for converting StatelessRandomUniformV2

这些警告是TensorFlow在将数据增强层的操作转换为计算图时的兼容性问题。数据增强用到的随机操作（如随机翻转、旋转、亮度调整依赖的随机数生成算子），在v2.9.1版本中没有对应的直接转换逻辑，框架只能退而用while_loop模拟这些操作的行为。这类警告属于框架内部的临时适配方案，不会影响代码功能，但会产生大量冗余输出。

有三种可行的处理方式：

1. 升级TensorFlow版本
   TensorFlow 2.10及以上版本针对这类算子的转换逻辑做了优化，已经修复了该警告问题。直接升级到最新稳定版本，即可彻底消除这些警告。

2. 针对性屏蔽警告
   如果暂时无法升级版本，可以通过过滤工具屏蔽特定警告：

   import tensorflow as tf
   # 精准屏蔽目标警告
   import warnings
   warnings.filterwarnings("ignore", category=UserWarning, message="Using a while_loop for converting.*")
   

   注意：若使用tf.get_logger().setLevel('ERROR')会全局屏蔽所有TensorFlow警告，若需保留其他重要警告，建议使用上述精准过滤方式。

3. 调整数据增强实现方式
   将数据增强层整合到模型结构中，而非通过dataset.map应用：

   def build_model(self):
       inputs = tf.keras.Input(shape=(img_height, img_width, 3))
       # 初始化数据增强模型
       aug_model = Sequential([
           layers.RandomFlip('horizontal_and_vertical'),
           layers.RandomRotation(30),
           layers.RandomTranslation(0.1, 0.1, fill_mode='nearest'),
           layers.RandomBrightness(0.1, value_range=(0.0, 1.0))
       ])
       x = aug_model(inputs, training=True)
       # 后续模型层定义
       x = layers.Conv2D(32, (3,3), activation='relu')(x)
       # ...其他层
       outputs = layers.Dense(num_classes, activation='softmax')(x)
       return tf.keras.Model(inputs, outputs)
   

   这种方式下，数据增强作为模型的一部分参与计算图构建，可避免map操作触发的算子转换警告，但需确保训练时training参数设为True，推理时设为False。

内容的提问来源于stack exchange，提问作者Karol Borkowski