当前位置：首页 > news >正文

Pytorch---ImageFolder

news 来源：原创 2025/6/2 11:56:18

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

基本概念与用途

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

root/dog/xxx.png
root/dog/xxy.png
root/dog/[...]/xxz.pngroot/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  # 自定义文件过滤函数
)

参数详解：

root
数据集的根目录路径，如 "/path/to/your/dataset"。
transform
对图像进行预处理的函数或变换序列，例如缩放、裁剪、归一化等。常用的变换包括 torchvision.transforms 中的类，如 Resize, ToTensor, Normalize。
target_transform
对标签（类别索引）进行预处理的函数，例如将标签转换为独热编码。
loader
自定义图像加载函数，默认使用 default_loader（基于 PIL）。可自定义以支持特殊格式（如 .tif）。
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 osdef 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 pathdataset = datasets.ImageFolder(root="/path/to/your/dataset",is_valid_file=is_valid_file
)

进阶技巧

1. 训练集与测试集分割

from torch.utils.data import random_splittrain_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)

注意事项

内存使用
图像数据通常较大，建议使用 DataLoader 的 batch_size 和 num_workers 参数优化内存和速度。
标签顺序
类别索引按字母顺序自动分配（如 ['cat', 'dog'] 对应 [0, 1]）。如需自定义顺序，可通过 class_to_idx 调整。
文件格式
默认支持 ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif']，其他格式需自定义 loader。
数据集清理
确保目录中没有无关文件，或使用 is_valid_file 过滤。

常见问题排查

问题1：加载失败或标签错误
- 原因：目录结构不符合要求或存在无关文件。
- 解决：检查目录结构，使用 dataset.classes 和 dataset.samples 验证。
问题2：内存溢出
- 解决：减小 batch_size，增加 num_workers，或使用内存映射文件。
问题3：图像损坏
- 解决：在加载时添加异常处理，或使用 is_valid_file 过滤损坏文件。