☰
基于PyTorch与CelebA的人脸识别模型训练实战解析
2026/10/11 12:21:01 网站建设 项目流程

简介:基于CelebA数据集与PyTorch实现的人脸识别神经网络项目,面向深度学习入门者及人脸识别方向开发者,覆盖数据加载、网络设计、模型训练与测试的完整流程,解决从零搭建人脸识别模型的实践问题。压缩包共30个文件,大小约24.72MB,包含10个Python脚本、6个XML/IML工程配置文件、6个TXT数据列表、4个pyc编译文件及1个pkl预训练权重;核心代码涵盖net_model.py网络定义、Dataset.py数据加载、train_model.py训练、test_model.py测试等模块,并附有README.md说明文档。资源结构清晰,便于按模块查阅和二次开发。目前已有116人在CSDN学习该资源。通过这份资源,读者既能了解CelebA数据集的属性标注与预处理方式,也能借助PyTorch动态图机制掌握神经网络训练的代码实现;基于现成的模型权重和脚本,可直接进行人脸检测或属性识别实验,或在此基础上开展迁移学习,节省大量调参和踩坑时间。

1. 这个压缩包里装的东西,为什么值得你花一个下午解开

做门禁机、考勤机或者安防摄像头的人脸识别模块时,很多工程师的第一反应是拿现成模型去跑推理,但真到要针对自己的场景做微调、要理解模型在干什么的时候,手里没一套能训练、能改、能复现的代码就会很被动。FaceDetectionByTorch.zip这个项目做的事情,就是基于CelebA(标题里拼成Celebra,是同一个数据集)提供的二十多万张人脸图像,用PyTorch从零训练一个人脸识别神经网络。它解决的不是“人脸检测”那个画框的问题,而是把人脸图像变成一串能区分身份的向量——这正是闸机对比、刷脸打卡背后真正在算的东西。

这个项目适合两类人。一类是刚把PyTorch环境折腾好、想找一个完整闭环练手的同学;另一类是已经在跑现成模型、但需要一套可修改的基线代码来做算法选型的人。它不追求刷榜,胜在结构完整:数据加载、模型定义、训练评估、权重保存都齐了。下面按我自己复现这类项目的习惯,从数据到训练再到排错,把这个方向一次讲透。

2. 先把数据吃透:CelebA的目录结构、标注格式和预处理管线

2.1 下载CelebA之前先想清楚:你到底要哪张表

CelebA全称是CelebFaces Attribute Dataset,由香港中文大学发布,包含10177个名人身份的202599张人脸图片,每张图都有40个二值属性标注(比如是否戴眼镜、是否微笑、是否金发),同时还有一个identity标注文件,记录每张图片属于哪个身份。对人脸识别任务来说,最有价值的是identity标注;对属性识别任务,才是那40个属性。这个项目标题落点是“人脸识别神经网络”,所以身份标注才是主菜。

下载的时候注意版本。CelebA官网提供的是ALign&Cropped版本,每张图已经按人脸关键点对齐并裁剪成178×218,这个版本最适合直接丢给CNN训练。原始版Img文件夹是整张照片,需要自己做检测和对齐,新手没必要碰。标注文件里最常用的是identity_CelebA.txt,每一行是“图片文件名 身份ID”,ID范围是1到10177。数据划分方面,官方给了train/val/test三个split文件,每个文件里是图片文件名列表。常见做法是直接用官方划分,省得自己分配的时候出现同一个人同时出现在训练集和验证集的泄漏问题。

提示:数据集解压后占用空间约1.5GB,标注文件是纯文本,几十KB到几百KB不等。训练前先确认磁盘余量,别把数据集和解压路径放在C盘系统盘上。

2.2 从原始图片到可训练张量:自定义Dataset的四个关键步骤

PyTorch训练人脸识别模型的第一步是写一个Dataset类。用torchvision自带的ImageFolder也行,但CelebA的标注是独立TXT文件,ImageFolder要求按文件夹组织类别,所以更稳的写法是自己解析标注文件。下面这段代码是我常用的CelebA Dataset实现核心部分,思路同样适用于其他“图片+独立TXT标注”的数据集。

import os from PIL import Image from torch.utils.data import Dataset class CelebADataset(Dataset): def __init__(self, img_dir, identity_file, split_file, transform=None): self.img_dir = img_dir self.transform = transform self.samples = [] # 每个元素是 (图片绝对路径, 身份ID) self._load(identity_file, split_file) def _load(self, identity_file, split_file): # 1. 读split文件拿到当前划分允许的图片名集合 allowed = set() with open(split_file, 'r') as f: for line in f: name = line.strip() if name: allowed.add(name) # 2. 读identity文件,只保留当前划分里的图片 with open(identity_file, 'r') as f: for line in f: parts = line.strip().split() if len(parts) != 2: continue img_name, identity = parts[0], int(parts[1]) if img_name in allowed: self.samples.append( (os.path.join(self.img_dir, img_name), identity) ) def __len__(self): return len(self.samples) def __getitem__(self, idx): img_path, label = self.samples[idx] image = Image.open(img_path).convert('RGB') if self.transform is not None: image = self.transform(image) return image, label

代码逻辑说明:_load方法先读split文件构建一个允许图片名的集合,再去扫identity文件逐行匹配。这样做的原因是训练集和验证集的图片不能混用,而identity文件包含全部20万张图,必须按划分过滤。两个文件都是逐行读,不用pandas也能跑,内存占用低。10万个样本的split文件大概几十毫秒就能读完。

参数说明:img_dir指向对齐裁剪后的图片目录;identity_file指向identity_CelebA.txt;split_file指向train.txt或val.txt;transform是后续要传入的数据增强和归一化。这个设计把划分逻辑和图片读取解耦,后面换验证集只需要换一个split文件,不需要改代码。

预处理里有两个细节容易被忽略。第一是图像要缩放到网络输入尺寸,CelebA原图是178×218,常见ResNet输入是224×224,直接resize会把脸拉变形,但CelebA已经对齐过人脸,轻微变形对识别影响不大,实践上直接resize是主流做法。第二是归一化必须用训练集统计的mean和std,ImageNet的均值[0.485, 0.456, 0.406]和标准差[0.229, 0.224, 0.225]在人脸数据上也基本适用,不用专门重算。数据增强方面,训练集建议加随机水平翻转(人脸左右对称,翻转不破坏身份信息),验证集只做resize和归一化,不加任何随机操作,保证评估可复现。

3. 网络结构怎么选:从骨干网络到特征输出头的落地搭配

3.1 为什么选ResNet当骨干而不是从零搭CNN

标题里写的是“神经网络”,但人脸识别领域真正能落地的骨干网络,常见选择是ResNet系列。原因很直接:人脸识别拼的是特征区分度,网络越深越能提取细粒度特征,但VGG这种纯堆卷积的结构在18层以上就会出现训练退化问题。ResNet引入残差连接(skip connection),让梯度可以跨层直接回传,训练更深网络不翻车。这个项目里如果用ResNet18做基线,参数量约1100万,单张2080Ti上训练一个batch 64张图、224×224输入,大概0.2秒一个step,属于性价比很高的起点。

选骨干还要看数据集规模。CelebA有1万多个身份,20多万张图,用ResNet18或ResNet34足够学到有区分度的特征;如果换成ResNet50,参数量翻几倍,训练时间线性拉长,但对CelebA这个量级的数据,精度提升有限。我的经验是先用ResNet18跑通全流程,确认loss和准确率都在合理区间后,再决定要不要换更大的骨干。

注意:如果换ResNet50,训练时间大约变成ResNet18的2.5到3倍。先小后大,避免一上来就把训练挂在那里空等。

3.2 分类头还是embedding输出:两种输出头的取舍

人脸识别网络的常见结构是“骨干网络 + 输出头”。输出头有两种做法。第一种是分类头:骨干网络后面接一个全连接层,输出维度等于身份数量(CelebA是10177类),用交叉熵损失训练。这种做法简单直接,模型收敛快,训练完把分类头最后一层之前的特征向量取出来,就是人脸embedding。第二种是度量学习头:输出一个固定维度的embedding(常见是128维或512维),用Triplet Loss或ArcFace训练,让同一个人脸向量距离近、不同人脸向量距离远。

这个项目标题强调的是“神经网络”整体实现,没有锁定具体损失。常见做法是先实现分类头+交叉熵,因为它对PyTorch新手最友好,不涉及样本挖掘这类复杂逻辑。ArcFace这类度量学习方法精度更高,但实现里有margin和scale两个超参要调,而且训练不稳定时会莫名掉点。我一般会把分类头版本作为主线代码跑通,同时预留embedding输出接口,为后面切换度量学习留后路。

import torch.nn as nn import torchvision.models as models class FaceNet(nn.Module): def __init__(self, num_classes=10177, embed_dim=512, dropout=0.3): super(FaceNet, self).__init__() # 加载预训练ResNet18,去掉最后一层全连接 backbone = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1) self.features = nn.Sequential(*list(backbone.children())[:-1]) # 池化后得到2048维特征,降维到embed_dim self.embed_fc = nn.Sequential( nn.Linear(2048, embed_dim), nn.BatchNorm1d(embed_dim), nn.ReLU(inplace=True), nn.Dropout(p=dropout) ) # 分类头,只在训练时用到 self.classifier = nn.Linear(embed_dim, num_classes) def forward(self, x, return_embedding=False): x = self.features(x) # 输出形状 (B, 2048, 1, 1) x = x.view(x.size(0), -1) # 展平成 (B, 2048) embed = self.embed_fc(x) # (B, embed_dim) if return_embedding: return embed logits = self.classifier(embed) # (B, num_classes) return logits

代码逻辑说明:先用torchvision自带的ResNet18预训练权重初始化主干,把最后两层(平均池化和全连接)拆掉,保留卷积部分输出2048维特征图。接一个自定义的embed_fc块,先降维到512维,再经过BatchNorm和ReLU。训练分类时走classifier输出10177维logits;提取特征时只需设置return_embedding=True,拿到的就是可直接做余弦相似度对比的512维向量。

参数说明:embed_dim设512是行业常见值,128维内存小但精度会略降,256维是折中;num_classes跟随数据集身份数,换数据集时这个参数必须同步改;dropout设0.3对防止过拟合有明显帮助。注意一个是pretrained参数在torchvision新版本里改成了weights接口,老代码用pretrained=True会告警,建议用weights=ResNet18_Weights.IMAGENET1K_V1这种写法。

直接用预训练权重还是随机初始化,这里有一个值得说透的细节。从零训练18层ResNet在CelebA上也能收敛,但需要更长轮数和更精细的学习率调度;加载ImageNet预训练权重后,模型已经学会了通用的边缘、纹理、颜色分布等底层特征,人脸数据集上只需要微调高层特征,收敛速度和最终精度都会更好。这个项目既然用的是PyTorch生态,非常建议加上预训练初始化,能避免很多“训练了好几个小时accuracy还是个位数”的尴尬。

4. 训练配置和评估协议:几个必调的参数与验证策略

4.1 训练脚本的关键参数:学习率、batch_size和优化器怎么配

训练人脸识别网络,最影响结果的是优化器、学习率策略和batch_size这三个旋钮。优化器方面我没有用纯SGD,而是选择了SGD加动量(momentum=0.9),这是ResNet系列最稳的搭配。Adam收敛快但容易停留在尖锐极小值,微调时用它没问题,但从头训练或从预训练继续训练,SGD的泛化能力明显更好。学习率初始值常见设置是0.01,配合CosineAnnealingLR在300个epoch内从0.01余弦退火到接近0,训练后期学习率小到能让loss在极小值附近稳定震荡。

batch_size直接受显卡显存限制。ResNet18输入224×224,batch_size=64时显存占用大约6GB到8GB,batch_size=128时接近12GB。如果你的显卡只有8GB显存,把batch_size降到32或48更安全。这里有一个很容易忽略的连带问题:batch_size变化后学习率也要跟着调。经验上,batch_size翻倍,初始学习率大概乘以1.5到2,否则batch大的模型更新步长会不稳定。我通常的做法是固定batch_size=64,初始lr=0.01,验证集acc上不去再调,而不是同时动两个变量。

import torch from torch import optim model = FaceNet(num_classes=10177).cuda() criterion = nn.CrossEntropyLoss() optimizer = optim.SGD( model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4 # L2正则化,抑制过拟合 ) scheduler = optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=300, # 总训练轮数 eta_min=1e-6 ) # 每轮训练核心循环(简化版) model.train() for images, labels in train_loader: images, labels = images.cuda(), labels.cuda() optimizer.zero_grad() logits = model(images) loss = criterion(logits, labels) loss.backward() optimizer.step() scheduler.step()

参数说明:weight_decay=5e-4是L2正则化的强度,太大(如1e-2)会让模型欠拟合,太小(如1e-5)起不到约束作用。5e-4是ResNet系列迁移学习里的常用默认值。T_max必须和总epoch数一致,否则退火曲线会中断。学习率调度器是在每个epoch结束后step(),而不是每个step都step,这个位置写错会导致学习率下降过快,训练后期几乎原地踏步。

训练时长要心里有数。CelebA训练集16万张图,batch_size=64,一个epoch是2500个step左右。ResNet18在单张RTX 3080上约0.15秒一步,一个epoch约6到7分钟,300个epoch就是30小时以上。实际不需要训满300轮,通常80到120轮时验证集准确率就开始进入平台期,可以提前停。显存不够又想训得快,可以开混合精度训练(AMP),PyTorch自带amp模块,能把训练时间压缩到原来的50%到60%,显存占用也明显下降。

4.2 别只看准确率:人脸识别的评估协议要讲清楚

分类任务的验证集top-1准确率能反映模型学到什么程度,但人脸识别的实际场景是“一对多比对”或“一对一确认”,单纯看top-1不够。常见做法是用验证集做两两比对测试:从验证集里随机抽1000对人脸,其中500对属于同一个人,500对属于不同人,计算每对的余弦相似度,设定一个阈值,然后统计在给定阈值下的准确率。这个准确率同时包含误拒率和误识率,更贴近门禁场景的真实效果。

实现时,先用训练好的模型把验证集所有图片过一遍,得到每张图的512维embedding并保存下来,然后随机组合正负样本对。这一步比训练还容易出问题,因为验证集里同一个人只有若干张图,正样本对数量有限。CelebA验证集同一个人通常有15到30张图,做2000对正样本没问题。下面这段代码给出了余弦相似度计算的骨架。

import torch import torch.nn.functional as F def compute_accuracy_at_threshold(embeddings, pairs, labels, threshold=0.5): # embeddings: (N, 512) 验证集所有图片的特征向量 # pairs: (M, 2) 每行是两张图的索引 # labels: (M,) 1表示同一个人,0表示不同人 emb1 = embeddings[pairs[:, 0]] # (M, 512) emb2 = embeddings[pairs[:, 1]] # (M, 512) cos_sim = F.cosine_similarity(emb1, emb2) # (M,) preds = (cos_sim >= threshold).long() correct = (preds == labels).sum().item() return correct / len(labels)

代码逻辑说明:cosine_similarity对每个向量对计算余弦相似度,归一化之后向量长度不影响结果,所以训练时embedding是否需要做L2归一化可以灵活处理。阈值threshold是后续要扫的参数:同一批验证对,分别用0.4、0.5、0.6去评估,画出准确率随阈值变化的曲线再选最优。验证集上最优的阈值,部署到真实闸机场景时还要再降一点,因为实际场景的图片质量远不如CelebA这种正脸、对齐过的数据。

评估时还有一个必须留意的细节:如果用了Dropout和BatchNorm,推理前必须把模型切到eval模式,否则每个batch的统计量不同,提取的embedding会抖动,直接影响相似度对比结果。很多新手在验证时用model.train()模式跑embedding,结果同一个人的两张图相似度只有0.6,排查半天发现是这个原因。

5. 避坑指南:训练人脸识别网络时最容易翻车的五个地方

5.1 PyTorch版本和CUDA不匹配,环境装好了模型也跑不起来

现象:执行model = FaceNet().cuda()时报错,提示CUDA driver version is insufficient或者找不到CUDA设备。很多人在Anaconda里装PyTorch时随手pip install torch,装到了CPU版本,后面所有GPU操作全部失效。

原因:PyTorch的CUDA支持分两层,一是PyTorch自带CUDA运行库,二是显卡驱动里的CUDA Driver。PyTorch版本要求的CUDA运行时版本不能高于驱动支持的版本。比如驱动只支持CUDA 11.8,你却装了cu121的PyTorch,就会报错。另一个坑是conda和pip混装导致torch被覆盖成CPU版。

解决:先跑nvidia-smi看右上角CUDA Version,再跑python -c "import torch; print(torch.__version__)"确认安装版本。安装时用conda install pytorch torchvision pytorch-cuda=11.8 -c pytorch这类显式指定CUDA版本的方式。一个快速自检命令是torch.cuda.is_available(),返回True才能走GPU训练。我一般装完立刻跑一个10行的最小网络做前向和反向,确认环境真的可用,而不是等到训练脚本跑到一半才发现问题。

5.2 loss变成NaN但模型没炸:学习率太高之外还有两个隐蔽原因

现象:训练前几个step loss正常,到某个step后突然变成NaN,之后一直NaN,模型输出的logits也出现inf。

原因:最常见是初始学习率太大导致梯度爆炸。但还有一个容易被忽略的情况——数据里有坏图。CelebA数据集通过网络下载,偶尔会有损坏的JPEG文件,Image.open()能打开但转成张量时出现异常像素值,前向传播后产生NaN梯度。第三个原因是未归一化的输入范围过大,0到255的像素值和预训练权重不匹配,损失函数在初始阶段就不稳定。

解决:训练脚本里把图片归一化到0到1范围(除以255),再用ImageNet的mean/std做标准化,这是标准做法但经常漏。加上梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0),能兜底梯度爆炸。数据加载时在__getitem__里用try-except捕获损坏图片,遇到坏图就返回同batch里另一张图。排查时先看loss在NaN之前那个step的梯度范数,如果梯度过大就是学习率问题;如果梯度正常,检查输入数据范围。

5.3 验证集准确率90%以上,闸机实测一用就翻车

现象:CelebA验证集top-1准确率达到92%,但放到真实门禁场景的人脸比对,识别率掉到60%到70%,误识别还特别多。

原因:这是典型的域偏移(domain shift)。CelebA全是明星写真、正脸、光线均匀、分辨率高;真实闸机场景有侧脸、逆光、运动模糊、戴帽子和口罩遮挡。模型在干净数据上学的特征,对遮挡和光照变化的鲁棒性不够。

解决:这是一个方向性判断,不是调参能解决的。如果目标是真实场景落地,训练数据必须掺入真实环境采集的人脸图,或者用数据增强模拟:随机亮度对比度扰动、随机遮挡(在图片上盖一块黑色矩形)、随机旋转正负10度、高斯模糊。数据增强不是加得越多越好,我之前做过一次实验,把增强强度拉满,模拟遮挡的区域占图片面积30%以上,验证集准确率反而掉了4个百分点,因为增强数据偏离真实分布太多。增强强度要逐步加,每次加完跑一个短训练看验证集趋势。

5.4 1万多个身份类别的分类头,过拟合和欠拟合同时存在

现象:训练集loss持续下降,验证集loss先降后升,典型的过拟合;但同一模型在部分高频身份上准确率很高,低频身份几乎全错。

原因:CelebA身份分布不均匀,头部身份有上百张图,尾部身份只有几张。分类头对尾部类别的权重更新不足,学不到有效特征。加上模型参数量大、训练数据类别多,整体容易过拟合。

解决:第一,数据增强用起来,尤其随机水平翻转,能翻倍有效样本;第二,加大weight_decay到1e-3试试,约束权重不过分增长;第三,采用类别均衡采样,每个batch尽量覆盖不同身份,避免某一轮被头部身份主导。实现上可以给Dataset加一个按照身份ID分组的索引,每次采样时先随机选身份,再从该身份里随机选一张图。这个改动对低频身份识别率的提升比调学习率明显得多。也可以给尾部身份做简单复制增强——虽然数据量增加不多,但让每个身份最少有8到10张图参与训练,分类边界更稳。

5.5 用ImageFolder偷懒加载CelebA,训练到一半发现数据对不上

现象:训练正常启动,但验证集准确率极低,抽查模型的预测结果,发现标签和图片对不上。

原因:ImageFolder要求每个类一个文件夹,CelebA下载下来不是这种结构——整个目录里是20万张平铺的图片,identity顺序也不是按文件夹组织的。如果自己按身份ID硬拆目录很容易出错,比如TXT里某一行是“000001.jpg 5”,把图片放进ID=5的文件夹时多一层或少一层路径,标签就全偏了。

解决:老老实实用自定义Dataset解析TXT,不要手动整理目录。一个极简的自检方法是加载训练集前100个样本,打印图片路径和标签,肉眼确认图片里的人和身份ID逻辑上能对上。还有就是训练集和验证集必须来自官方split文件,别自己随机划分,否则同一个人可能同时出现在训练集和验证集,验证指标虚高,模型实际泛化能力仍然不行。

6. 让模型真正可用:特征向量比对、阈值标定和ONNX导出

训练好模型只是第一步,把模型变成能对外提供服务的东西才算完成。最后一章说两个我常用的验证和落地技巧。

第一个是embedding的余弦相似度可视化验证。训练结束后,不要只看准确率,把几个人的测试图片过一遍模型,提取512维向量,算两人之间的余弦相似度矩阵。同一人的相似度应该在0.6到0.9区间,不同人的相似度在0.1到0.4区间。如果同一个人的相似度比不同人还低,说明特征没学好,回去加训练轮数或调学习率。这个矩阵用matplotlib画热力图画出来,比任何指标都直观。顺便可以做一个简单的t-SNE降维可视化,看不同身份的embedding在二维平面上是否分开,快速判断特征空间的质量。

第二个是导出ONNX,让人脸识别模型能跑到C++服务端或边缘设备的推理引擎上。PyTorch模型查权重文件离线加载,灵活性不够,ONNX是跨平台通用格式,OpenCV的DNN模块、TensorRT、OpenVINO都能直接读。导出前先把模型切到eval模式,固定输入尺寸,然后用torch.onnx.export导出。

import torch import onnxruntime as ort model = load_trained_model() # 加载训练好的权重 model.eval() dummy_input = torch.randn(1, 3, 224, 224).cuda() torch.onnx.export( model, dummy_input, "face_recognition.onnx", input_names=["input"], output_names=["embedding"], dynamic_axes=None, # 固定batch_size=1,简单稳定 opset_version=12 ) # 用onnxruntime验证导出结果是否一致 ort_session = ort.InferenceSession("face_recognition.onnx") ort_out = ort_session.run(None, {"input": dummy_input.cpu().numpy()})[0] torch_out = model(dummy_input, return_embedding=True).detach().cpu().numpy() print("最大误差:", torch_out - ort_out) # 应该接近1e-5以内

代码逻辑说明:dynamic_axes=None意味着导出的模型只支持batch_size=1的输入,这在闸机场景够用——每次刷脸就是单张图过模型。如果要支持批量比对(比如一帧画面同时识别多张脸),可以把dynamic_axes设置成{"input": {0: "batch"}},但某些推理引擎对动态shape支持不好,建议固定尺寸保平安。opset_version=12是兼容性和新算子特性的平衡点。

导出后误差验证这一步不能省。PyTorch和ONNX Runtime在BatchNorm、Reshape等算子的实现细节上可能有微小差异,最大误差超过1e-3就说明导出过程有问题,通常出在自定义算子或过于老旧的opset版本。

我自己的习惯是训练完一版模型,先做一轮余弦相似度抽查,确认特征空间没问题,再导ONNX跑一次端到端推理,最后才考虑部署。这套流程走下来,踩过的坑大多集中在数据加载和训练配置上,模型结构本身反而不太容易出错。希望这篇笔记能帮你把这个项目从“解压即跑”变成“知道每一行在干什么、出了问题知道去哪查”。

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询