Keras中基于含大量类别的DataFrame图像训练方案咨询

完全理解你的困境——类别太多没法分目录，数据又大到装不下内存，确实不能直接用flow_from_directory或者简单的flow。这里有几个非常实用的方案，都是Keras生态里的标准解法：

1. 首选：用ImageDataGenerator.flow_from_dataframe #

这其实是Keras专门为你的场景设计的API！它直接支持从DataFrame读取图像路径和标签，不需要把图像按类别分目录，而且是批量加载数据，不会一次性把所有图像塞进内存。

举个简单的使用示例：

from tensorflow.keras.preprocessing.image import ImageDataGenerator

# 初始化数据生成器，可加入预处理/数据增强逻辑
datagen = ImageDataGenerator(rescale=1./255, validation_split=0.2)

# 训练集生成器
train_generator = datagen.flow_from_dataframe(
    dataframe=your_df,
    x_col='image_path',  # 你的DataFrame中存储图像路径的列名
    y_col='class_label', # 存储类别标签的列名
    target_size=(224, 224), # 输入图像的统一尺寸
    batch_size=32,
    class_mode='categorical', # 多分类用这个，单分类用'binary'，整数标签用'int'
    subset='training'
)

# 验证集生成器
val_generator = datagen.flow_from_dataframe(
    dataframe=your_df,
    x_col='image_path',
    y_col='class_label',
    target_size=(224, 224),
    batch_size=32,
    class_mode='categorical',
    subset='validation'
)

# 直接用生成器训练模型
model.fit(train_generator, validation_data=val_generator, epochs=10)

这个方法会自动处理标签编码（比如把字符串标签转成one-hot向量），还支持验证集拆分、数据增强，完全适配你的需求。

2. 更灵活：自定义Sequence生成器 #

如果你的预处理逻辑比较复杂（比如需要自定义图像加载、特殊的数据增强），可以继承keras.utils.Sequence实现自己的数据生成器。这个类是Keras官方推荐的，支持多进程训练，而且线程安全。

示例代码框架：

from tensorflow.keras.utils import Sequence
import cv2
import numpy as np

class CustomImageGenerator(Sequence):
    def __init__(self, df, img_size, batch_size, preprocess_func=None):
        self.df = df
        self.img_size = img_size
        self.batch_size = batch_size
        self.preprocess_func = preprocess_func
        # 建立标签到整数的映射
        self.classes = df['class_label'].unique()
        self.label_map = {cls: idx for idx, cls in enumerate(self.classes)}

    def __len__(self):
        # 返回总批次数
        return np.ceil(len(self.df) / self.batch_size).astype(int)

    def __getitem__(self, idx):
        # 获取当前批次的样本
        batch_df = self.df.iloc[idx*self.batch_size : (idx+1)*self.batch_size]
        batch_images = []
        batch_labels = []

        for _, row in batch_df.iterrows():
            # 加载图像（也可以用PIL等其他库）
            img = cv2.imread(row['image_path'])
            img = cv2.resize(img, self.img_size)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 转成RGB格式
            # 自定义预处理
            if self.preprocess_func:
                img = self.preprocess_func(img)
            batch_images.append(img)
            # 处理标签
            label = self.label_map[row['class_label']]
            batch_labels.append(label)

        # 转成numpy数组，标签转one-hot（按需选择）
        batch_images = np.array(batch_images)
        batch_labels = np.eye(len(self.classes))[batch_labels]
        return batch_images, batch_labels

# 使用自定义生成器
train_gen = CustomImageGenerator(your_train_df, (224,224), 32)
val_gen = CustomImageGenerator(your_val_df, (224,224), 32)
model.fit(train_gen, validation_data=val_gen, epochs=10)

3. 高性能选择：用tf.data.Dataset #

如果用的是TensorFlow 2.x的Keras，tf.data.Dataset是更底层、性能更高的方案，适合超大规模数据集。它支持异步加载、预取，能充分利用硬件资源。

示例代码：

import tensorflow as tf

def load_image(image_path, label):
    # 加载图像
    img = tf.io.read_file(image_path)
    img = tf.image.decode_jpeg(img, channels=3)
    img = tf.image.resize(img, (224, 224))
    # 预处理（示例用ResNet的预处理逻辑）
    img = tf.keras.applications.resnet50.preprocess_input(img)
    # 标签编码（如果是字符串标签）
    label = tf.argmax(tf.equal(label, tf.constant(your_df['class_label'].unique())), axis=0)
    return img, label

# 从DataFrame创建Dataset
ds = tf.data.Dataset.from_tensor_slices(
    (your_df['image_path'].values, your_df['class_label'].values)
)

# 映射加载函数、批量、预取
ds = ds.map(load_image, num_parallel_calls=tf.data.AUTOTUNE)
ds = ds.batch(32)
ds = ds.prefetch(tf.data.AUTOTUNE)

# 拆分训练/验证集
train_size = int(0.8 * len(your_df))
train_ds = ds.take(train_size)
val_ds = ds.skip(train_size)

# 训练模型
model.fit(train_ds, validation_data=val_ds, epochs=10)

小提醒 #

• 不管用哪种方法，确保图像路径是绝对路径，或者相对路径相对于当前工作目录，否则会出现读取错误。
• 如果你的标签是整数形式，class_mode可以设为'int'，模型最后用SparseCategoricalCrossentropy损失函数会更高效，不用手动转one-hot。
• 自定义生成器时，记得加异常捕获（比如图像损坏无法读取），避免训练中断。

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