使用image_dataset_from_directory遇TensorFlow不可分片数据集报错求助

使用image_dataset_from_directory构建数据集，结合MirroredStrategy多GPU训练CNN时，出现以下警告并导致脚本终止：

W tensorflow/core/grappler/optimizers/data/auto_shard.cc:786] AUTO sharding policy will apply DATA sharding policy as it failed to apply FILE sharding policy because of the following reason: Found an unshardable source dataset: name: "TensorSliceDataset/_1"
op: "TensorSliceDataset"

训练代码如下：

train_ds = image_dataset_from_directory(
    self.data_dir,
    validation_split=self.val,
    subset="training",
    seed=1,
    image_size=(self.height, self.width),
    batch_size=BATCH_SIZE)

val_ds = image_dataset_from_directory(
    self.data_dir,
    validation_split=self.val,
    subset="validation",
    seed=1,
    image_size=(self.height, self.width),
    batch_size=BATCH_SIZE)

gpus = tf.config.list_logical_devices('GPU')
strategy = tf.distribute.MirroredStrategy(gpus)
with strategy.scope():
    model = Sequential([
        Conv2D(32, 3, padding='same'),
        PReLU(),
        MaxPooling2D(),
        Flatten(),
        Dense(64),
        PReLU(),
        Dense(num_classes)])
    optimizer = Adam()

    model.compile(optimizer=optimizer,
                  loss=SparseCategoricalCrossentropy(from_logits=True),
                  metrics=['accuracy'])

history = model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=100,
)

方案1：针对image_dataset_from_directory直接调整分片策略 #

这个警告的核心原因是：使用validation_split时，image_dataset_from_directory内部会通过TensorSliceDataset切分训练/验证集，导致无法应用FILE分片策略，只能自动降级为DATA分片。显式指定分片策略即可解决问题，同时不影响原有数据集构建流程：

# 创建train_ds和val_ds之后添加以下代码
options = tf.data.Options()
# 显式指定使用DATA分片策略
options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA
train_ds = train_ds.with_options(options)
val_ds = val_ds.with_options(options)

# 可选：添加缓存和预取优化多GPU训练效率
train_ds = train_ds.cache().prefetch(tf.data.AUTOTUNE)
val_ds = val_ds.cache().prefetch(tf.data.AUTOTUNE)

# 后续的策略定义、模型编译和训练代码保持不变

方案2：手动构建数据集（使用from_tensor_slices处理多类别目录） #

如果希望彻底避免这类分片问题，可以手动遍历多类别目录生成文件路径和标签，再用from_tensor_slices构建数据集，步骤如下：

import os
import sklearn.model_selection

# 1. 遍历多类别目录，收集文件路径和对应标签
data_dir = self.data_dir
class_names = sorted(os.listdir(data_dir))
class_to_idx = {name: idx for idx, name in enumerate(class_names)}

file_paths = []
labels = []
for cls_name in class_names:
    cls_dir = os.path.join(data_dir, cls_name)
    # 获取当前类别下所有图片文件
    cls_files = tf.io.gfile.glob(os.path.join(cls_dir, "*"))
    file_paths.extend(cls_files)
    labels.extend([class_to_idx[cls_name]] * len(cls_files))

# 2. 手动划分训练和验证集
train_paths, val_paths, train_labels, val_labels = sklearn.model_selection.train_test_split(
    file_paths, labels, test_size=self.val, random_state=1
)

# 3. 定义图片加载和预处理函数
def load_and_preprocess(path, label):
    img = tf.io.read_file(path)
    # 根据图片格式调整解码方式（如decode_png）
    img = tf.image.decode_jpeg(img, channels=3)
    img = tf.image.resize(img, (self.height, self.width))
    # 可添加其他预处理（如归一化）
    return img, label

# 4. 构建数据集并优化
train_ds = tf.data.Dataset.from_tensor_slices((train_paths, train_labels))
train_ds = train_ds.map(load_and_preprocess, num_parallel_calls=tf.data.AUTOTUNE)
train_ds = train_ds.shuffle(buffer_size=len(train_paths)).batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)

val_ds = tf.data.Dataset.from_tensor_slices((val_paths, val_labels))
val_ds = val_ds.map(load_and_preprocess, num_parallel_calls=tf.data.AUTOTUNE)
val_ds = val_ds.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)

# 后续的策略定义、模型编译和训练代码保持不变

这种方式完全由你控制数据集的构建流程，生成的from_tensor_slices数据集可以更好地适配多GPU的分片策略，避免自动分片的警告问题。

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