图像分类：创建高效的数据集生成器

本文将介绍如何使用 Keras 和 PyTorch 创建高效的数据集生成器，用于图像分类任务。生成器用于将图像数据加载到模型中进行训练、验证和测试。

Keras 数据集生成器

train_generator = train_datagen.flow_from_directory(
    train_dir,
    target_size=(150, 150),
    batch_size=20,
    class_mode='categorical'
)

validation_generator = validation_datagen.flow_from_directory(
    validation_dir,
    target_size=(150, 150),
    batch_size=20,
    class_mode='categorical'
)

test_generator = validation_datagen.flow_from_directory(
    test_dir,
    target_size=(150, 150),
    batch_size=20,
    class_mode='categorical'
)

PyTorch 数据集生成器

train_dataset = datasets.ImageFolder(train_dir, transform=train_transforms)
train_generator = torch.utils.data.DataLoader(train_dataset, batch_size=20, shuffle=True)

validation_dataset = datasets.ImageFolder(validation_dir, transform=validation_transforms)
validation_generator = torch.utils.data.DataLoader(validation_dataset, batch_size=20, shuffle=True)

test_dataset = datasets.ImageFolder(test_dir, transform=test_transforms)
test_generator = torch.utils.data.DataLoader(test_dataset, batch_size=20, shuffle=True)

代码说明

• train_dir, validation_dir, test_dir: 存放训练、验证和测试图像数据的目录路径
• target_size: 将图像调整为的目标尺寸
• batch_size: 每次迭代中加载的图像数量
• class_mode: 指定标签的类型（'categorical' 表示 one-hot 编码）
• train_transforms, validation_transforms, test_transforms: 用于数据增强的变换

总结

通过使用这些示例代码，您可以轻松创建有效的数据集生成器，以训练、验证和测试您的图像分类模型。这将有助于提高模型的性能和泛化能力。