TensorFlow使用Dataset.map调用RandomRotation时报错的问题求解

• 该问题为TensorFlow 2.4~2.6版本的已知bug，仅出现在tf.data.Dataset.map中调用旧版RandomRotation（位于layers.experimental.preprocessing下）的场景
• RandomRotation的底层实现需要动态生成旋转矩阵，旧版本实现中隐含了依赖numpy类型转换的逻辑：而Dataset.map默认在Graph模式下运行，所有张量都是无实际值的符号张量，无法完成numpy数组转换，因此触发报错
• RandomFlip、RandomContrast的实现逻辑更轻量，没有涉及numpy转换的步骤，因此可以在map中正常运行
• 把增强层嵌入网络内部时，Keras会自动完成层的Graph追踪，适配符号张量的计算逻辑，因此不会触发报错

方案1：升级TensorFlow版本 #

直接升级到TensorFlow 2.7及以上版本，该版本已修复此问题，RandomRotation可直接在Dataset.map中正常调用。

方案2：使用tf.py_function包裹增强逻辑 #

强制增强函数在Eager模式下运行，规避符号张量的转换限制，代码示例如下：

import tensorflow as tf

data_augmentation = tf.keras.Sequential([
    layers.experimental.preprocessing.RandomRotation(0.2)
])

def augmentation(image,label):
    image = data_augmentation(image)
    return image, label

def aug_wrapper(image, label):
    # 注意Tout的类型要和你的实际数据类型匹配
    img, lab = tf.py_function(func=augmentation, inp=[image, label], Tout=[tf.float32, tf.int32])
    # 手动恢复张量形状，避免后续训练时形状丢失
    img.set_shape(image.shape)
    lab.set_shape(label.shape)
    return img, lab

ds_train = ds_train.cache()
ds_train = ds_train.map(aug_wrapper, num_parallel_calls=AUTOTUNE)

方案3：替换为原生API实现随机旋转 #

直接使用tensorflow_addons的旋转接口实现相同逻辑，完全规避Keras层的兼容性问题：

import tensorflow as tf
import tensorflow_addons as tfa
import math

def augmentation(image, label):
    # 0.2对应RandomRotation的参数，即最大旋转角度为±0.2*2π
    rotate_angle = tf.random.uniform(shape=[], minval=-0.2*2*math.pi, maxval=0.2*2*math.pi)
    # fill_mode可按需调整为constant/nearest等，和RandomRotation默认行为保持一致
    image = tfa.image.rotate(image, rotate_angle, fill_mode='reflect')
    return image, label

ds_train = ds_train.cache()
ds_train = ds_train.map(augmentation, num_parallel_calls=AUTOTUNE)

方案4：保持增强层嵌入模型的写法 #

这也是TensorFlow官方推荐的数据增强实现方式：将增强层直接接在模型输入层之后，不仅可以规避上述bug，还可以保证模型推理时自动包含增强逻辑（如需），跨设备部署的兼容性也更好。

内容的提问来源于stack exchange，提问作者Oskar Zdrojewski