Pytorch---ImageFolder

torchvision.datasets.ImageFolder 是 PyTorch 中用于加载图像数据集的实用类，特别适合处理按文件夹组织的图像数据。

基本概念与用途

ImageFolder 假设数据集按照以下结构组织：

root/dog/xxx.png
root/dog/xxy.png
root/dog/[...]/xxz.png

root/cat/123.png
root/cat/nsdf3.png
root/cat/[...]/asd932_.png

其中，每个子文件夹（如 dog, cat）代表一个类别，文件夹名称即为类别名，文件夹内包含该类别的图像。这种结构在图像分类任务中非常常见。

核心参数

torchvision.datasets.ImageFolder(
    root: str,                      # 数据集根目录
    transform: Optional[Callable] = None,  # 图像预处理转换
    target_transform: Optional[Callable] = None,  # 标签预处理转换
    loader: Callable[[str], Any] = default_loader,  # 自定义图像加载函数
    is_valid_file: Optional[Callable[[str], bool]] = None  # 自定义文件过滤函数
)

参数详解：

1. root

   数据集的根目录路径，如 "/path/to/your/dataset"。

2. transform

   对图像进行预处理的函数或变换序列，例如缩放、裁剪、归一化等。常用的变换包括 torchvision.transforms 中的类，如 Resize, ToTensor, Normalize。

3. target_transform

   对标签（类别索引）进行预处理的函数，例如将标签转换为独热编码。

4. loader

   自定义图像加载函数，默认使用 default_loader（基于 PIL）。可自定义以支持特殊格式（如 .tif）。

5. is_valid_file

   自定义文件过滤函数，用于决定哪些文件应该被加载。返回 True 表示文件有效。

核心属性与方法

1. classes

返回所有类别的名称列表，例如 ['cat', 'dog']。

2. class_to_idx

返回类别名称到索引的映射字典，例如 {'cat': 0, 'dog': 1}。

3. samples

返回所有样本的元组列表，格式为 (file_path, class_index)，例如：

[('/path/to/cat/123.png', 0), ('/path/to/dog/xxx.png', 1), ...]

4. __len__

返回数据集的样本总数。

5. __getitem__

通过索引获取单个样本 (image, label)，其中：

• image 是经过 transform 处理后的图像张量。
• label 是经过 target_transform 处理后的类别索引。

使用示例

1. 基础用法

from torchvision import datasets, transforms

# 定义图像预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载数据集
dataset = datasets.ImageFolder(root="/path/to/your/dataset", transform=transform)

# 创建数据加载器
from torch.utils.data import DataLoader
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 遍历数据
for images, labels in dataloader:
    print(f"Batch shape: {images.shape}, Labels: {labels}")

2. 自定义图像加载器

处理特殊格式（如 .tif）时：

from PIL import Image
import os

def tif_loader(path):
    return Image.open(path).convert('RGB')

dataset = datasets.ImageFolder(
    root="/path/to/your/dataset",
    loader=tif_loader,
    extensions=('.tif', '.tiff')
)

3. 自定义文件过滤

只加载特定文件：

def is_valid_file(path):
    return path.endswith('.jpg') and 'train' in path

dataset = datasets.ImageFolder(
    root="/path/to/your/dataset",
    is_valid_file=is_valid_file
)

进阶技巧

1. 训练集与测试集分割

from torch.utils.data import random_split

train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size
train_dataset, test_dataset = random_split(dataset, [train_size, test_size])

2. 多进程数据加载

使用 num_workers 参数加速数据读取：

dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)

3. 不同的预处理策略

对训练集和测试集应用不同的变换：

train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

test_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

train_dataset = datasets.ImageFolder(root="/path/to/train", transform=train_transform)
test_dataset = datasets.ImageFolder(root="/path/to/test", transform=test_transform)

注意事项

1. 内存使用

   图像数据通常较大，建议使用 DataLoader 的 batch_size 和 num_workers 参数优化内存和速度。

2. 标签顺序

   类别索引按字母顺序自动分配（如 ['cat', 'dog'] 对应 [0, 1]）。如需自定义顺序，可通过 class_to_idx 调整。

3. 文件格式

   默认支持 ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif']，其他格式需自定义 loader。

4. 数据集清理

   确保目录中没有无关文件，或使用 is_valid_file 过滤。

常见问题排查

• 问题1：加载失败或标签错误

  • 原因：目录结构不符合要求或存在无关文件。
  • 解决 ：检查目录结构，使用 dataset.classes 和 dataset.samples 验证。

• 问题2：内存溢出

  • 解决 ：减小 batch_size，增加 num_workers，或使用内存映射文件。

• 问题3：图像损坏

  • 解决 ：在加载时添加异常处理，或使用 is_valid_file 过滤损坏文件。

总结

ImageFolder 是 PyTorch 中处理图像分类数据的强大工具，通过简单的目录结构即可自动构建数据集。核心优势在于：

• 自动处理类别标签映射。
• 灵活的预处理和加载机制。
• 与 DataLoader 无缝集成，支持批量加载和多进程加速。

掌握 ImageFolder 后，你可以轻松处理各种图像分类任务，如猫狗识别、花卉分类等。