在 Tensorflow 中，使用 Keras 预处理层进行数据增强时出现形状错误

Question

我正在训练一个 CNN，并通过以下定义通过 Keras 层应用数据增强：

data_augmentation = keras.Sequential([
  layers.experimental.preprocessing.RandomRotation(factor=0.4, fill_mode="wrap"),
  layers.experimental.preprocessing.RandomTranslation(height_factor=0.2, width_factor=0.2, fill_mode="wrap"),
  layers.experimental.preprocessing.RandomFlip("horizontal"),
  layers.experimental.preprocessing.RandomContrast(factor=0.2),
  layers.experimental.preprocessing.RandomHeight(factor=0.2),
  layers.experimental.preprocessing.RandomWidth(factor=0.2)
])

以下是相关代码的 sn-ps：

def process(x, y):
  x = DATASETS_DIR + "/" + x + ".jpg"
  x = tf.io.read_file(x)
  x = tf.image.decode_jpeg(x, channels=3)
  x = tf.image.resize(x, [299, 299])
  x = layers.experimental.preprocessing.Rescaling(1./255)(x)
  
  return x, y

def process_with_augmentation(x, y):
  x, y = process(x, y)

  x = data_augmentation(x)

  return x, y

train_ds = train_ds.map(process_with_augmentation, num_parallel_calls=tf.data.experimental.AUTOTUNE)
validation_ds = validation_ds.map(process, num_parallel_calls=tf.data.experimental.AUTOTUNE)

如果我在process_with_augmentation() 中注释掉x = data_augmentation(x)，则代码可以正常工作。如果我不注释掉，我会收到以下错误：

ValueError: Shape must be rank 4 but is rank 3 for '{{node sequential/random_rotation/transform/ImageProjectiveTransformV3}} = ImageProjectiveTransformV3[dtype=DT_FLOAT, fill_mode="WRAP", interpolation="BILINEAR"](rescaling/add, sequential/random_rotation/rotation_matrix/concat, sequential/random_rotation/transform/strided_slice, sequential/random_rotation/transform/fill_value)' with input shapes: [299,299,3], [299,8], [2], [].

关于如何解决这个问题的任何想法？

Answer 1

keras 层需要批量大小。当您将函数映射到 tf.data.Dataset 时，图像将缺少该批次维度。您可以通过在调用 keras 预处理模型之前添加维度来解决此问题：

def process_with_augmentation(x, y):
  x, y = process(x, y)

  x = tf.expand_dims(x, axis=0) # adding the batch dimension
  x = data_augmentation(x)
  x = tf.squeeze(x, [0]) # removing the batch dimension
  return x, y

另一种选择是在调用模型之前调用tf.data.Dataset.batch，即：

train_ds = train_ds.batch(BATCH_SIZE)
train_ds = train_ds.map(lambda x,y:(data_augmentation(x),y))

如果你想摆脱那个额外的维度，你可以在增强数据集上调用unbatch。

---

不过，这不是解决这个问题的最优雅的方法。我建议直接将您的预处理模型集成到您的神经网络中，或者使用 tf.image 模块中的函数重写您的预处理逻辑。