PyTorch加载CelebA数据集全程指南：自定义Dataset与DataLoader避坑手册-平芜编程栈

1. 项目概述：CelebA数据集与PyTorch图片加载的典型问题

很多人第一次接触CelebA数据集，是在做人脸属性识别、人脸关键点检测或者生成对抗网络实验的时候。这个数据集的全称是CelebFaces Attributes Dataset，由某研究机构发布，里面包含了超过20万张名人面部图片，每张图都有40个属性标注，比如是否戴眼镜、是否微笑、头发颜色、性别等。因为量大、标注全、图片规格统一，它几乎成了人脸相关深度学习实验的标准数据集之一。但问题在于，这个数据集在学术社区里被广泛使用，真正把它下载下来、用PyTorch加载进模型训练时，各种意想不到的坑就接踵而至。

我在做某个图像分类Demo的时候，第一次尝试用这个数据集训练一个人脸属性分类模型。原本以为把图片下载下来、用torchvision.datasets.ImageFolder加载就行了，结果从文件目录结构到数据预处理、再到显存分配，一路踩坑踩到怀疑人生。后来我花了好几天时间，把PyTorch加载本地图片文件夹的整套流程彻底捋了一遍，才弄明白问题到底出在哪些环节。这篇文章就结合CelebA这个具体数据集，把我实际遇到的坑和对应解决方案完整记录下来。

适合看这篇文章的读者有两类，一类是刚接触CelebA数据集、正准备跑人脸相关实验的学生或开发者，一类是已经在用PyTorch加载本地图片文件夹、但总在数据加载环节出各种奇怪问题的人。看完之后你至少能搞清楚三件事：第一，CelebA数据集应该怎么正确准备和放置目录；第二，ImageFolder和DataLoader背后的运行逻辑到底是什么；第三，遇到加载报错时，怎么快速定位问题根源而不是靠乱猜。

2. 数据集的本质理解：CelebA到底解决了什么问题

2.1 数据集的规模与标注体系

CelebA数据集的核心价值，在于它提供了大规模且带丰富标注的人脸图片。总共有202599张图片，每张都是3通道的彩色图。原始图片的分辨率并不完全一致，但经过对齐裁剪后，标准版本会统一到178像素宽、218像素高。这个尺寸在深度学习里不算大，但对人脸识别和属性分析任务来说足够了。

更关键的是标注部分。每一张图片都有40个二值属性标注，比如：

性别（Male）
是否微笑（Smiling）
是否戴眼镜（Eyeglasses）
头发颜色（Black_Hair、Blond_Hair、Brown_Hair等）
是否年轻（Young）
是否浓妆（Heavy_Makeup）
是否戴帽子（Hat）
是否有双下巴（Double_Chin）

此外数据集还提供了5个关键点坐标标注（左眼、右眼、鼻子、左嘴角、右嘴角）以及人脸检测框。这意味着同一个数据集可以同时支撑属性分类、关键点检测、人脸验证、人脸生成等多个方向的任务，这也是它成为"万金油"数据集的原因。

2.2 数据集目录结构与组织方式

CelebA数据集的压缩包解压后，主要包含以下几个部分：

celeba/ ├── img_align_celeba/ # 对齐后的人脸图片文件夹 │ ├── 000001.jpg │ ├── 000002.jpg │ └── ... ├── list_attr_celeba.txt # 属性标注文件，40列二值标注 ├── list_bbox_celeba.txt # 人脸框坐标标注 ├── list_landmarks_align_celeba.txt # 关键点坐标标注 ├── list_eval_partition.txt # 训练/验证/测试集的划分文件 └── identity_CelebA.txt # 身份标注

图片文件夹里清一色是000001.jpg这类以数字命名的文件，没有子文件夹，也没有按类别分目录。这个细节非常重要，因为torchvision.datasets.ImageFolder默认要求目录按照"每个类别一个子文件夹"的格式组织。很多人第一次用CelebA时都会栽在这一点上，直接拿ImageFolder去加载，结果发现一个类别都没有，报错说Found no valid file或者Number of classes is 0。

2.3 为什么很多人栽在目录格式上

ImageFolder的工作逻辑其实很死板：根目录下的每个子文件夹会被视为一个类别，子文件夹里的图片就是这个类别的样本。比如你建了root/dog/和root/cat/两个文件夹，训练时模型就会自动把dog映射到标签0、cat映射到标签1。

但CelebA的图片都在一个平铺的文件夹里，没有子目录，所以直接用ImageFolder根本行不通。就算你强行指定root=img_align_celeba，它也会因为没有子目录而报错。我在第一次实验时就是这个情况，当时以为数据集损坏了，重新下载了一遍，结果还是一样。

这个问题的本质是：ImageFolder是给"目录即标签"的场景设计的，而CelebA的标签存在list_attr_celeba.txt文本文件里，属于"外部标注文件"的场景。想用ImageFolder处理CelebA，必须先按属性把图片分到不同文件夹里，或者干脆不用ImageFolder，自定义一个Dataset类来读取标注文件。

3. PyTorch加载图片文件夹的常规方案与原理

3.1 ImageFolder + DataLoader的标准流程

在没有遇到CelebA之前，很多人包括我在内，用ImageFolder加载图片数据集已经形成了一套固定写法：

from torchvision import datasets, transforms from torch.utils.data import DataLoader transform = transforms.Compose([ transforms.Resize((128, 128)), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) dataset = datasets.ImageFolder(root='path/to/root', transform=transform) dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)

这段代码背后做了三件事：

第一，扫描根目录下所有子文件夹，给每个子文件夹分配一个从0开始的类别索引，存储在dataset.classes里。

第二，遍历所有图片文件，读取文件路径，存储一个(图片路径, 类别索引)的映射表。

第三，在训练过程中，DataLoader按批次从映射表中取出样本，把图片路径传给transform里的PIL.Image.open()读取，然后做缩放、转张量、归一化等一系列操作。

这个过程看起来简单，但问题恰恰出在"自动扫描"和"自动分配"两个环节上。ImageFolder只能识别有限几种图片格式，比如.jpg、.jpeg、.png、.bmp、.gif，如果文件夹里混入了其他格式文件，尤其是那些没有扩展名的文件，它就会自动跳过。更坑的是，如果文件夹里混入了损坏的图片文件，ImageFolder能"看到"它，但后期读取时就会直接报错。

3.2 dataset.py中自定义Dataset的必要性

对于CelebA这种标签在外部文本文件里的数据集，正解是自定义一个Dataset子类来接管数据读取逻辑。Dataset类是PyTorch所有数据加载逻辑的基类，只需要实现两个方法：

__len__()：返回数据集总样本数
__getitem__(idx)：返回第idx个样本的（图片tensor，标签tensor）对

自定义类的核心优势在于灵活性。你可以随意定义图片路径的解析方式、标签的读取方式，甚至可以动态做数据增强。相比ImageFolder的"目录即标签"逻辑，自定义Dataset让"文件路径和标签完全分离"成为可能。

我在实际操作中写了一个专门处理CelebA属性的CelebADataset类，大致思路是：

读取list_attr_celeba.txt，把40个属性按行存成列表。
根据list_eval_partition.txt判断哪些图片属于训练集、验证集或测试集。
在__getitem__里，用图片路径去img_align_celeba目录下找对应jpg文件并读取。

这样就不用把图片按属性重新复制到不同文件夹里，省了至少几十GB的磁盘空间和大量等待时间。

3.3 图像读取、变换与张量转换链条

无论用ImageFolder还是自定义Dataset，最终图片都要经过同样的读取处理链：

文件路径 → PIL Image对象 → transform变换 → torch.FloatTensor

transforms库的作用就是把一张PIL图片转换成网络需要的格式。最基础的三个转换是：

Resize((H, W))：把图片缩放到指定尺寸，Image.BICUBIC是默认插值方式，一般够用。
ToTensor()：把PIL Image从H×W×C的uint8数组变成C×H×W的float32张量，数值范围从[0,255]缩放到[0,1]。
Normalize(mean, std)：按通道做标准化，公式是(x - mean) / std，处理后的数据近似服从标准正态分布，有利于模型收敛。

这三个变换的顺序不能乱。ToTensor必须放在Resize后面，因为Resize操作的是PIL对象，而Normalize必须放在ToTensor后面，因为它操作的是tensor。我见过有人把Normalize放在ToTensor前面，结果报AttributeError: 'PIL.Image.Image' object has no attribute 'std'。

4. 实操过程：从CelebA下载到PyTorch成功加载的全流程

4.1 下载与文件检查

CelebA数据集本身的下载地址在很多镜像站都有，下载时注意选择对齐版本即可。下载完成后，先做几件基础检查工作，避免后面踩坑：

检查文件完整性。压缩包大概1.3GB左右，解压后图片文件夹大约16GB，如果解压过程中报错中断，很可能是下载文件不完整，需要重新下载。

检查图片文件夹结构。确认img_align_celeba文件夹下至少有一张jpg图片，比如000001.jpg。如果解压后文件夹是空的，多半是解压工具的问题，建议换用专门的解压工具再试一次。

检查标注文件格式。用文本编辑器打开list_attr_celeba.txt，第一行是"image_id"，第二行是40个属性名，第三行开始才是具体数据。这个格式后面解析时很重要。

4.2 目录结构调整的两种正确思路

如果你坚持用ImageFolder，那么必须调整目录结构。调整方式有两种：

第一种方式是按单个属性分目录。比如想训练微笑分类器，就把图片按Smiling=1和Smiling=-1分成两个文件夹。这种方式简单直接，适合属性较少、只跑单个任务的实验。

第二种方式是按身份分目录。CelebA有10177个不同的人，按identity_CelebA.txt里的身份编号建子文件夹，每个文件夹放一个人的所有图片。这种方式适合做人脸验证或识别任务。

但无论哪种方式，都需要写脚本遍历大约20万张图片并复制或移动文件。我做过一次，整台电脑跑了一个多小时才完成，磁盘占用翻倍。从那之后我再也不这么干了，改用自定义Dataset，直接读外部标注文件。

4.3 自定义CelebA Dataset类详解

自定义Dataset是绕开目录结构限制最省事的方法。我基于list_attr_celeba.txt写了一个可用的类，代码大致如下：

import os from PIL import Image import torch from torch.utils.data import Dataset class CelebADataset(Dataset): def __init__(self, img_dir, attr_file, transform=None): self.img_dir = img_dir self.transform = transform self.img_names = [] self.attrs = [] with open(attr_file, 'r') as f: lines = f.readlines() # 第一行是"image_id"，第二行是属性名，第三行开始是数据 self.attr_names = lines[1].strip().split() for line in lines[2:]: parts = line.strip().split() img_name = parts[0] values = [int(v) for v in parts[1:]] self.img_names.append(img_name) self.attrs.append(values) def __len__(self): return len(self.img_names) def __getitem__(self, idx): img_path = os.path.join(self.img_dir, self.img_names[idx]) image = Image.open(img_path).convert('RGB') if self.transform: image = self.transform(image) attr = torch.tensor(self.attrs[idx], dtype=torch.float32) return image, attr

这里有几个值得注意的细节：

Image.open()默认返回PNG或JPEG的原始模式，有时可能是灰度图或RGBA图。.convert('RGB')强制统一成三通道，避免后续训练时通道数不匹配。

属性标注文件里的值是1和-1，表示正向和负向属性。但PyTorch的BCEWithLogitsLoss期望标签范围是0到1，所以读取时要提前处理好标签转换。通常做法是(values + 1) // 2，把-1映射成0、1映射成1。

torch.tensor在Python 3.9以上版本可以直接使用，旧版需要torch.FloatTensor。建议直接用torch.tensor(..., dtype=torch.float32)，这是官方推荐的写法。

4.4 用DataLoader批量加载与性能调优

有了自定义Dataset，关联DataLoader就顺理成章：

train_dataset = CelebADataset( img_dir='./img_align_celeba', attr_file='./list_attr_celeba.txt', transform=transform ) train_loader = DataLoader( train_dataset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True )

num_workers参数值得单独说明。它决定用几个子进程加载数据。设成0表示在主进程里加载，简单但慢。设成4或8表示开多个子进程并行加载，效率高很多。

但num_workers也不是越大越好。Windows系统下，超过8往往会导致内存暴涨，甚至直接卡死。另外子进程太多时，每个进程都会复制一份数据集索引，内存开销成倍增加。实测下来4到6是一个甜点区间。

pin_memory=True在做GPU训练时很有用。它把数据放在锁页内存里，GPU拷贝数据的效率会明显提升。如果你用的是CPU训练，这个参数基本没用。

4.5 训练循环中常见的显存与数据维度陷阱

模型训练时最常见的两个报错，我在CelebA上也都碰过。

第一个是维度不匹配。ImageFolder默认返回单张图片的tensor形状是(C, H, W)，但DataLoader在堆叠成batch时会自动变成(B, C, H, W)。如果你的数据集类在__getitem__里多添加了一个batch维度，堆出来的形状就会变成(B, 1, C, H, W)，模型第一层卷积直接报错说期望4维输入但得到5维。

第二个是标签类型不匹配。BCEWithLogitsLoss要求标签是float32类型，而且形状要跟预测一致。很多人习惯用to(device)把输入和标签搬到GPU，却忘了标签要.float()转换，结果报Expected object of scalar type Long but got scalar type Float之类的错误。

我在处理CelebA属性标签时，标签矩阵是40维的浮点向量，正好对应模型输出的40个logits，用BCEWithLogitsLoss一次性计算所有属性的损失。这个设计思路是：每个属性互不排斥，一个人可以同时是男性且戴眼镜且微笑，所以用sigmoid加二值交叉熵，而不是CrossEntropyLoss。

5. 常见问题速查与排查实录

5.1 加载失败类的错误集合

错误现象	根本原因	解决方案
`Found no valid file in path`	文件夹内没有子目录，ImageFolder无法识别类	改用自定义Dataset，或先按类分目录
`RuntimeError: Found 0 files in subfolders`	图片格式不被ImageFolder支持	检查图片扩展名，统一转成jpg或png
`PIL.UnidentifiedImageError`	图片文件损坏或格式伪装	写脚本扫描所有图片，删除或替换损坏文件
`EOFError: Ran out of input`	图片在读取时被截断	重新下载该图片，或加载时加异常处理跳过
`FileNotFoundError`	路径拼接错误，文件名对不上	打印实际路径检查，确认文件名前缀和补零位数

5.2 内存与性能类的错误集合

错误现象	根本原因	解决方案
加载数据时内存爆炸	num_workers设太高，子进程复制了过多数据	调低num_workers，限制到4以内
训练速度极慢	每次iter都在做IO读取和图片解码	把图片缓存到内存或使用内存映射文件
显存不足OOM	batch_size太大或图片分辨率太高	降低batch_size，或采用混合精度训练
DataLoader假死	Windows上的多进程问题	把数据加载代码放进`if __name__ == '__main__'`保护块

5.3 我自己踩过的三个坑

第一个坑是文件名补零问题。CelebA的图片命名是固定8位数字，比如000001.jpg。但有一次我写代码生成路径时用了str(i).zfill(6)，结果文件名变成了000001.jpg（6位），前缀补零不够导致找不到文件。这个问题排查了很久才发现，因为报错提示FileNotFoundError，我一开始以为是目录路径写错了，完全没想到是文件名本身的问题。

第二个坑是裁剪尺寸不一致引发的通道报错。CelebA原始图片是178×218，但我在Resize里写了transforms.Resize(128)，这只会缩放短边到128，长边按比例缩放，结果不同图片的最终尺寸不一致，堆叠batch时报错说尺寸不匹配。后来我改成transforms.Resize((128, 128))，强制宽高都缩放到128，问题才解决。

第三个坑是DataLoader迭代到一半突然卡死。排查后发现是Windows环境下多进程加载数据时触发了递归调用问题。解决办法是在脚本入口添加if __name__ == '__main__':保护，把训练逻辑全包进去，之后再也没有出现过假死。

5.4 调参与加速的几点独家心得

关于图片加载的加速，我最终尝试了一套组合方案，效果很好。先把所有图片转换成RGB模式并统一缩放到合适尺寸，保存为压缩率适中的jpg文件，避免读取时大量IO开销。然后num_workers设为4，pin_memory设为True，是不是GPU训练都有稳定提升。另外把torch.cuda.amp.autocast()打开，在精度影响可控的前提下，显存占用可以降低三分之一，batch_size可以适度调大。

如果你想做更极致的加速，可以尝试把所有图片一次性读入内存。20万张178×218×3的图片，每张大约100KB到200KB，全部读入内存大概需要20GB到30GB。如果机器内存够大，这样做能彻底消除IO瓶颈，训练迭代速度会快非常多。内存不够的话，可以用torchdata或webdataset这类工具做流式加载，原理类似，但内存友好得多。

6. 从加载到业务：一个完整的人脸属性分类Demo

6.1 数据集划分逻辑

CelebA官方提供了list_eval_partition.txt，用数字0、1、2分别表示训练集、验证集和测试集。这个划分文件遵循的原则是：同一个人不会同时出现在训练集和验证集里，避免了数据泄露问题。

我在自定义Dataset时，通过读取这个文件来筛选需要的子集。比如训练模型时只保留标记为0的图片，验证时只保留标记为1的图片。这样写的好处是，整个数据集只需要一份，不同阶段用同一个实例的不同子集就行。

6.2 模型选型与损失函数

人脸属性分类是个多标签分类问题。我选的模型结构很简单：用torchvision.models.resnet18做主干网络，把最后一层全连接改成输出40维的线性层。

import torch.nn as nn from torchvision import models model = models.resnet18(pretrained=True) in_features = model.fc.in_features model.fc = nn.Linear(in_features, 40)

损失函数用BCEWithLogitsLoss，原因是CelebA的40个属性之间不是互斥关系，一个人可以同时有多个属性。如果用CrossEntropyLoss，就被迫把40个二分类问题压缩成一个多分类问题，信息损失非常大。我在实验中发现，BCEWithLogitsLoss配合sigmoid输出，模型收敛速度和准确率都比多分类方案好得多。

6.3 训练过程中的监控技巧

训练时任选几个batch，每一轮都打印一下标签分布和预测分布。我曾经遇到过一种情况，模型预测值几乎全是负数，sigmoid输出全接近0，代表所有属性都预测为不存在。排查后发现问题出在标签处理上：原始标签是1和-1，我直接丢给BCEWithLogitsLoss，而它期望的标签范围是0到1，导致损失异常大，模型被带偏。

改成(label + 1) / 2之后，数值范围变成0到1，训练恢复正常。这个细节在list_attr_celeba.txt的说明文档里完全没有提到，完全要靠自己踩坑才会发现。

另一个监控指标是属性级别的准确率，不要只盯着整体准确率。CelebA里有些属性非常不平衡，比如"秃顶"（Bald）在全部图片中占比很低，模型即使全部预测为"不秃顶"，整体准确率也还能看，但这个属性本身的准确率就是0。这种时候应该单独计算每个属性的F1分数，才能发现模型是不是在偷懒。

7. 结语：加载图片文件夹并不是简单事

从准备CelebA数据集到最终能把图片正确喂给PyTorch模型，中间涉及目录结构、标注解析、张量变换、多进程加载、显存管理等一系列环节。任何一环出了问题，整个训练流程都会卡住。坦白说，这些坑在官方文档里很难一次性找到完整解释，很多时候要靠一步一步打印日志、检查形状、验证数值才能定位到根本原因。

我个人的体会是，处理图片数据集的加载时，思路要转变一下：先搞清楚数据结构长什么样，再决定用什么方式读取。CelebA的标注在外部文本文件里，就不应该硬套ImageFolder的"目录即标签"逻辑；文件名定了就按文件名匹配，就不要随意改补零位数；模型期望什么标签分布，就提前把标签转换做好。想清楚这些，踩坑的概率会大幅度降低。

如果你正在用自己的图片文件夹数据集，建议从一开始就把自定义Dataset作为首选方案，虽然代码量多一点，但灵活性高出的不是一点半点。特别是你后面如果要加数据增强、做多模态输入、调整负样本比例，都会省力很多。

最后再分享一个实用习惯：给自定义Dataset加上调试模式，通过参数控制是否打印每个样本的图片路径和标签值。数据加载出问题时，这行打印信息往往能直接告诉你问题出在哪一层，省去大量猜测时间。这个习惯我保留至今，帮我节省了非常多排查时间。