CelebA数据集高效处理实战：从数据清洗到模型训练的完整解决方案

1次阅读

没有评论

共计 1558 个字符，预计需要花费 4 分钟才能阅读完成。

CelebA 数据集是人脸识别领域的标杆数据集，包含超过 20 万张名人图像，每张图像标注了 40 多种面部属性（如是否戴眼镜、是否微笑等）。虽然数据丰富，但在实际使用中常遇到以下问题：

内存占用高：一次性加载所有图像会导致内存爆炸（约 200GB+）
属性对齐复杂：多标签标注文件需要特殊处理才能与图像正确匹配
预处理效率低：传统单线程处理方式耗时过长（原始分辨率图像处理需 10+ 小时）

常见处理方案主要有两种：

全量加载方案
优点：实现简单，随机访问快
缺点：内存需求极高，不适合普通开发机
流式处理方案
优点：内存友好
缺点：I/ O 频繁导致速度慢

本文采用 分块预处理 + 内存映射 的混合方案：

将数据集按 8:1:1 比例分块（训练 / 验证 / 测试）
使用 HDF5 格式存储预处理后的图像和标注
通过内存映射实现零拷贝读取

使用 PIL 库的懒加载特性，配合多进程加速：

from PIL import Image
from multiprocessing import Pool

def process_image(img_path):
    with Image.open(img_path) as img:
        return img.resize((128, 128)).convert('RGB')

with Pool(8) as p:  # 8 进程并行
    processed_imgs = p.map(process_image, img_paths)

将 CSV 标注转换为 HDF5 格式，提升检索速度：

import h5py

with h5py.File('celeba_attrs.h5', 'w') as f:
    # 存储为二维数组 [样本数×属性数]
    f.create_dataset('attributes', data=attr_matrix, compression='gzip')

import torch
from torch.utils.data import Dataset

class CelebADataset(Dataset):
    def __init__(self, h5_path):
        self.h5 = h5py.File(h5_path, 'r', swmr=True)

    def __getitem__(self, idx):
        # 内存映射读取，避免全量加载
        img = torch.from_numpy(self.h5['images'][idx])
        attrs = torch.from_numpy(self.h5['attributes'][idx])
        return img, attrs