TensorFlow中for循环内修改数据集不生效问题排查

我尝试通过在for循环中随机调整图像色相来增强数据集，但修改无法在循环外持久化。使用tf.keras.utils.image_dataset_from_directory导入数据集，代码如下：

def augment(image, label, counter):
  randNr = tf.random.uniform(shape=(), minval=-1, maxval=1, dtype=tf.dtypes.float32)
  image = tf.image.adjust_hue(image, delta=randNr)
  #display some values
  if(counter<1):
    print(randNr)
    plt.figure()
    plt.imshow(image[0].numpy().astype("uint8"))
    plt.show()
  return image, label

temp1 = 0
for image, label in v_dataset:
  image, label = augment(image, label, temp1)
  #display some values
  if(temp1<1):
    plt.figure()
    plt.imshow(image[0].numpy().astype("uint8"))
    plt.show()
  temp1 += 1

#display some values
plt.figure(figsize=(10, 10))
for images, labels in v_dataset.take(1):
    print("images shape: ", np.shape(images))
    for i in range(9):
        ax = plt.subplot(3, 3, i + 1)
        plt.imshow(images[i].numpy().astype("uint8"))
        plt.title(int(labels[i]))
        plt.axis("off")
plt.show()

现象：前两次打印图像时，色相已按预期变化（如图1、图2），但后续从原数据集取图像打印时，所有图像均无色相变化（如图3）。

原因分析 #

• TensorFlow的Dataset是惰性计算的，你在for循环里修改的image只是当前批次的临时张量副本，根本没改动原始的v_dataset。每次遍历原数据集，它都会重新从数据源加载原始数据，之前的临时修改不会被保留。
• 循环里的操作相当于“读出来改一改，但没存回去”，原数据集自然还是原样。

解决方法 #

用TensorFlow官方推荐的Dataset.map()方法把增强逻辑绑定到数据集上，这样每次遍历数据集都会自动应用增强，修改会持久化到新的数据集对象中：

import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np

def augment(image, label):
    # 去掉counter参数，不需要依赖外部计数
    randNr = tf.random.uniform(shape=(), minval=-1, maxval=1, dtype=tf.dtypes.float32)
    image = tf.image.adjust_hue(image, delta=randNr)
    return image, label

# 生成带增强的新数据集，原v_dataset保持不变
augmented_v_dataset = v_dataset.map(augment)

# 验证增强效果
temp1 = 0
for image, label in augmented_v_dataset:
    if(temp1 < 1):
        plt.figure()
        plt.imshow(image[0].numpy().astype("uint8"))
        plt.show()
    temp1 += 1

# 验证循环外的增强效果
plt.figure(figsize=(10, 10))
for images, labels in augmented_v_dataset.take(1):
    print("images shape: ", np.shape(images))
    for i in range(9):
        ax = plt.subplot(3, 3, i + 1)
        plt.imshow(images[i].numpy().astype("uint8"))
        plt.title(int(labels[i]))
        plt.axis("off")
plt.show()

补充说明 #

• map()会遍历数据集的每个元素并应用增强函数，返回全新的增强数据集，原数据集不受影响。
• 每次遍历增强后的数据集时，都会重新生成随机色相值，保证增强的随机性，符合数据集增强的需求。
• 如果需要固定增强结果（比如测试阶段），可以把随机操作改成固定值，或者单独写一个无随机的增强函数。

内容的提问来源于stack exchange，提问作者Rumpel Stilzchen