RetinaFace模型训练全流程:从数据标注到模型部署
如果你对人脸检测感兴趣,想自己动手训练一个能精准识别人脸和关键点的模型,那这篇文章就是为你准备的。今天咱们不聊那些复杂的理论,直接上手,从数据准备开始,一步步带你完成RetinaFace模型的训练,直到最后把它部署起来,变成一个能用的工具。
RetinaFace这个模型挺有意思的,它不仅能框出人脸在哪,还能顺带找出眼睛、鼻子、嘴角这些关键点,算是一举两得。整个过程听起来可能有点复杂,但别担心,我会尽量用大白话把每个环节讲清楚,你跟着做就行。
1. 训练前的准备工作:环境和数据
在开始敲代码之前,咱们得先把“厨房”收拾好,也就是准备好开发环境和训练数据。
1.1 搭建Python开发环境
我建议你直接用Anaconda来管理环境,这样能避免很多包版本冲突的麻烦事。打开你的终端或者命令提示符,跟着下面的命令一步步来。
首先,创建一个新的Python环境,这里我们用Python 3.8,比较稳定:
conda create -n retinaface_train python=3.8 -y创建好后,激活这个环境:
conda activate retinaface_train接下来,安装我们训练模型需要的核心库。PyTorch是必须的,你可以根据自己电脑有没有GPU来选择安装命令。如果有NVIDIA的显卡,建议安装GPU版本,训练会快很多。
# 如果你有GPU(CUDA 11.3为例),安装这个 pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113 # 如果只有CPU,安装这个 # pip install torch==1.12.1+cpu torchvision==0.13.1+cpu torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cpu然后安装其他必要的工具包:
pip install opencv-python pillow numpy matplotlib scipy tqdm pip install tensorboard # 用来可视化训练过程环境到这里就基本配好了。
1.2 获取并认识WiderFace数据集
模型训练得好不好,数据是关键。RetinaFace论文里用的就是WiderFace数据集,咱们也用这个。它包含了各种场景、尺度、姿态和遮挡的人脸,非常丰富。
你可以去WiderFace的官网下载数据集。下载后,它的目录结构通常是这样的:
WIDER_FACE/ ├── WIDER_train/ │ ├── images/ │ │ ├── 0--Parade/ │ │ ├── 1--Handshaking/ │ │ └── ... (其他61个场景文件夹) │ └── annotations/ │ └── 这里存放着对应的标注文件 ├── WIDER_val/ │ ├── images/ │ └── annotations/ └── wider_face_split/ ├── wider_face_train_bbx_gt.txt ├── wider_face_val_bbx_gt.txt └── ...重点要关注的是wider_face_split文件夹里的那几个.txt文件。它们不是常见的XML或JSON格式,而是一种特定的文本格式,记录了每张图片里人脸框的位置和关键点信息。咱们下一步就要处理它。
2. 数据处理的实战环节
拿到原始数据后,不能直接扔给模型,需要先“加工”一下。
2.1 把标注文件转换成模型认识的格式
WiderFace提供的标注文件需要被解析,并转换成模型训练时需要的格式。通常,我们会把每张图片的标注信息(人脸框坐标、关键点坐标)存成一个单独的.txt文件,或者整合到一个大的JSON文件里。
下面是一个简单的Python脚本,用来解析原始的标注文件。你需要把path_to_labels改成你自己wider_face_split文件夹的路径。
import os import cv2 def parse_wider_face_annotation(annotation_path, images_dir): """ 解析WiderFace的标注文件。 annotation_path: 例如 'wider_face_split/wider_face_train_bbx_gt.txt' images_dir: 图像根目录,例如 'WIDER_train/images' """ with open(annotation_path, 'r') as f: lines = f.readlines() data = [] i = 0 while i < len(lines): # 第一行是图片文件名 image_name = lines[i].strip() i += 1 # 第二行是人脸数量 num_faces = int(lines[i].strip()) i += 1 image_path = os.path.join(images_dir, image_name) # 检查图片是否存在(可选) if not os.path.exists(image_path): # 跳过不存在的图片,并消耗掉对应的bbox行 i += num_faces continue bboxes = [] landmarks = [] for _ in range(num_faces): # 读取一个人脸的标注:bbox (x1, y1, w, h, blur, expression, illumination, invalid, occlusion, pose) bbox_info = list(map(int, lines[i].strip().split()[:4])) # 只取前4个:x1, y1, w, h i += 1 # WiderFace原始标注没有提供5点关键点,所以这里先置空。 # 如果你的数据有关键点,需要在这里读取。 # 例如:landmark_info = list(map(float, lines[i].strip().split()[:10])) # 5个点,10个值 # i += 1 # landmarks.append(landmark_info) landmarks.append([0]*10) # 占位符 # 将 (x1, y1, w, h) 转换为 (x1, y1, x2, y2) x1, y1, w, h = bbox_info bboxes.append([x1, y1, x1+w, y1+h]) # 存储这张图片的信息 img = cv2.imread(image_path) if img is None: continue height, width = img.shape[:2] data.append({ 'image_path': image_path, 'width': width, 'height': height, 'bboxes': bboxes, 'landmarks': landmarks }) return data # 使用示例 train_data = parse_wider_face_annotation('path_to_labels/wider_face_train_bbx_gt.txt', 'WIDER_train/images') print(f"成功加载 {len(train_data)} 张训练图片")重要提示:标准的WiderFace数据集只提供人脸框,不提供五官关键点。如果你想训练带关键点检测的RetinaFace,需要自己标注关键点,或者使用其他包含关键点的数据集(如CelebA)。上面的代码中,landmarks部分我用零做了占位,如果你有自己的关键点数据,需要修改解析逻辑。
2.2 让数据“增强”,模型更健壮
如果只给模型看原始图片,它很容易“死记硬背”,遇到没见过的角度、亮度就懵了。数据增强就是通过一些随机变换,人工制造出更多样的训练样本。
咱们在训练时,可以实时地对每批图片进行增强。这里我用albumentations这个强大的库,先安装它:pip install albumentations。
然后,定义一个增强管道:
import albumentations as A from albumentations.pytorch import ToTensorV2 def get_train_transform(): return A.Compose([ # 随机水平翻转,人脸对称,这个增强很有效 A.HorizontalFlip(p=0.5), # 随机调整亮度、对比度,模拟不同光照 A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.5), # 随机缩放裁剪,让模型适应不同大小的人脸 A.RandomSizedBBoxSafeCrop(height=640, width=640, erosion_rate=0.2, p=0.5), # 标准化图像像素值,并转换为Tensor A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ToTensorV2(), ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['labels'])) # 注意格式 # 注意:上面的变换包含了BBox的变换,使用时需要确保标注信息(bboxes)也一同传入变换函数。这个管道做了几件事:随机左右翻转图片、改变亮度和对比度、随机裁剪出一块640x640的区域。注意,裁剪时erosion_rate参数可以防止把人脸框裁掉太多。最后把图片归一化并转成PyTorch需要的Tensor格式。
3. 模型训练的核心步骤
数据处理完后,就进入重头戏——训练模型。这里咱们不从头造轮子,基于一个开源的RetinaFace PyTorch实现来修改和训练。
3.1 构建数据加载器
首先,我们需要创建一个PyTorch的Dataset类来加载我们处理好的数据。
from torch.utils.data import Dataset, DataLoader import torch import cv2 import numpy as np class WiderFaceDataset(Dataset): def __init__(self, data_list, transform=None): self.data_list = data_list self.transform = transform def __len__(self): return len(self.data_list) def __getitem__(self, idx): item = self.data_list[idx] image_path = item['image_path'] bboxes = item['bboxes'] # 每个bbox格式: [x1, y1, x2, y2] # 读取图片 image = cv2.imread(image_path) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 转成RGB # 为每个bbox生成一个标签(这里简单处理,人脸都为1) labels = [1] * len(bboxes) # 应用数据增强 if self.transform: # albumentations需要bbox是列表格式,且每个bbox末尾带一个类别标签 transformed = self.transform(image=image, bboxes=bboxes, labels=labels) image = transformed['image'] bboxes = transformed['bboxes'] # 增强后bbox坐标可能变了 labels = transformed['labels'] # 将bboxes和labels转换为Tensor # 注意:这里需要根据你的模型输入要求,进一步处理成目标格式(如先验框匹配)。 # 以下是一个简化的返回,实际训练代码中需要更复杂的匹配过程。 targets = { 'boxes': torch.tensor(bboxes, dtype=torch.float32) if bboxes else torch.zeros((0, 4), dtype=torch.float32), 'labels': torch.tensor(labels, dtype=torch.int64), } return image, targets # 创建数据加载器 train_dataset = WiderFaceDataset(train_data, transform=get_train_transform()) train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, collate_fn=lambda x: tuple(zip(*x))) # 注意collate_fn这里有个关键点,collate_fn函数。因为每张图片里的人脸数量不一样,导致targets的尺寸不统一,无法直接堆叠成一个Batch。上面的写法是PyTorch检测任务中常见的处理方式,它返回的是一个列表,列表里每个元素是(image, target)对。
3.2 理解并配置损失函数
RetinaFace的损失函数是“多任务”的,可以把它想象成同时要完成好几份作业:
- 分类损失:判断一个锚点框(Anchor)里装的是不是人脸。
- 边界框回归损失:如果是人脸,这个框的位置准不准,需要微调。
- 关键点回归损失:人脸关键点的位置准不准。
在开源代码中,这部分通常已经实现好了,叫做MultiBoxLoss。你需要关注的是几个损失项的权重,它们决定了模型更“看重”哪个任务。在配置文件中,你可能会看到类似这样的参数:
# 在配置中(如 cfg.py 或 config.py 中) cfg['loss_weights'] = { 'cls': 1.0, # 分类损失权重 'box': 0.5, # 边界框回归损失权重 'landmark': 0.1, # 关键点回归损失权重 }调参小建议:刚开始训练时,可以先把landmark的权重设得低一点(比如0.01),甚至设为0(如果你没有关键点标注)。让模型先学会找到人脸,再慢慢学定位关键点。如果一开始权重太高,模型可能因为关键点任务太难而“学歪了”。
3.3 启动训练与监控
假设我们已经有了模型定义 (model)、优化器 (optimizer) 和损失函数 (criterion),训练循环的大致框架如下:
import torch from torch.utils.tensorboard import SummaryWriter from tqdm import tqdm device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') model.to(device) model.train() writer = SummaryWriter('logs/retinaface') # 用于TensorBoard可视化 num_epochs = 100 for epoch in range(num_epochs): epoch_loss = 0.0 loop = tqdm(train_loader, desc=f'Epoch [{epoch+1}/{num_epochs}]') for batch_idx, (images, targets) in enumerate(loop): images = [img.to(device) for img in images] # 需要根据你的模型结构,将targets处理成合适的格式并转移到device # targets = ... 处理过程 ... optimizer.zero_grad() # 前向传播:模型会输出分类、框回归、关键点的预测值 preds = model(images) # 计算多任务损失 loss_cls, loss_box, loss_landmark = criterion(preds, targets) total_loss = loss_cls + loss_box + loss_landmark # 反向传播和优化 total_loss.backward() optimizer.step() epoch_loss += total_loss.item() loop.set_postfix(loss=total_loss.item()) # 记录到TensorBoard writer.add_scalar('Loss/total', total_loss.item(), epoch * len(train_loader) + batch_idx) writer.add_scalars('Loss/components', { 'cls': loss_cls.item(), 'box': loss_box.item(), 'landmark': loss_landmark.item() }, epoch * len(train_loader) + batch_idx) avg_loss = epoch_loss / len(train_loader) print(f'Epoch {epoch+1} 平均损失: {avg_loss:.4f}') # 每10个epoch保存一次模型 if (epoch + 1) % 10 == 0: torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': avg_loss, }, f'checkpoints/retinaface_epoch_{epoch+1}.pth') writer.close()训练时,多观察TensorBoard里的损失曲线。理想情况下,总损失和各个分项损失都应该随着训练逐渐下降并趋于平稳。如果发现损失剧烈震荡或者不下降,可能是学习率太高、数据有问题或者模型结构不合适。
4. 模型部署与使用
训练完成后,得到一个.pth权重文件。怎么把它用起来呢?
4.1 加载模型进行推理
部署的第一步是写一个推理脚本,加载训练好的权重,并对新图片进行预测。
import torch import cv2 import numpy as np from model import RetinaFace # 假设你的模型定义在这个文件里 from utils import PriorBox, decode, decode_landmark, py_cpu_nms # 需要从开源代码中导入这些后处理函数 def load_model(checkpoint_path, device='cuda'): """加载训练好的模型""" # 初始化模型结构(必须和训练时一模一样) cfg = {...} # 需要填入训练时使用的配置 model = RetinaFace(cfg=cfg, phase='test') model.to(device) model.eval() # 切换到评估模式 # 加载权重 checkpoint = torch.load(checkpoint_path, map_location=device) model.load_state_dict(checkpoint['model_state_dict']) print(f"模型从 {checkpoint_path} 加载成功") return model, cfg def predict_image(model, cfg, image_path, device='cuda', confidence_threshold=0.5): """对单张图片进行预测""" # 预处理图片 img_raw = cv2.imread(image_path, cv2.IMREAD_COLOR) img = cv2.cvtColor(img_raw, cv2.COLOR_BGR2RGB) im_height, im_width, _ = img.shape # 将图片缩放并归一化,具体缩放比例需参考训练配置 scale = torch.Tensor([im_width, im_height, im_width, im_height]) img = cv2.resize(img, (cfg['image_size'], cfg['image_size'])) img = img.astype(np.float32) img -= (104, 117, 123) # RetinaFace常用的均值减去 img = img.transpose(2, 0, 1) # HWC -> CHW img = torch.from_numpy(img).unsqueeze(0).to(device) # 前向推理 with torch.no_grad(): loc, conf, landms = model(img) # 后处理:解码先验框,应用NMS非极大值抑制 # 这里需要调用原项目中的 decode, decode_landmark, py_cpu_nms 等函数 # 因为代码较长,以下为伪代码逻辑 priors = PriorBox(cfg).forward() # 生成先验框 boxes = decode(loc.data.squeeze(0), priors, cfg['variance']) # 解码得到真实框 boxes = boxes * scale # 缩放回原图尺寸 scores = conf.squeeze(0)[:, 1] # 取人脸的置信度 # 解码关键点(如果有) if landms is not None: landms = decode_landmark(landms.data.squeeze(0), priors, cfg['variance']) landms = landms * scale.repeat(1, 5) # 5个关键点 # 应用置信度阈值和NMS inds = np.where(scores.cpu() > confidence_threshold)[0] boxes = boxes[inds] scores = scores[inds] if landms is not None: landms = landms[inds] # NMS过滤重叠框 dets = torch.cat([boxes, scores.unsqueeze(1)], dim=1) if boxes.shape[0] > 0 else torch.Tensor([]) keep = py_cpu_nms(dets.cpu().numpy(), nms_threshold=0.4) dets = dets[keep] final_landms = landms[keep] if landms is not None else None # 在原图上画结果 for b in dets: x1, y1, x2, y2, score = b cv2.rectangle(img_raw, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2) cv2.putText(img_raw, f'{score:.2f}', (int(x1), int(y1)-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2) if final_landms is not None: for lm in final_landms: lm = lm.reshape(-1, 2) for (x, y) in lm: cv2.circle(img_raw, (int(x), int(y)), 2, (0, 0, 255), -1) return img_raw, dets, final_landms # 使用示例 device = 'cuda' if torch.cuda.is_available() else 'cpu' model, cfg = load_model('checkpoints/retinaface_epoch_100.pth', device) result_img, boxes, landmarks = predict_image(model, cfg, 'your_test_image.jpg', device) cv2.imwrite('result.jpg', result_img) print(f"检测到 {len(boxes)} 张人脸")4.2 转换为ONNX格式以便跨平台部署
如果你想在移动端、边缘设备或者不同的深度学习框架中使用这个模型,最好把它转换成ONNX格式。
import torch torch.onnx.export( model, # 模型 dummy_input, # 模型输入(一个示例Tensor) "retinaface.onnx", # 输出文件名 export_params=True, # 导出权重 opset_version=11, # ONNX算子集版本 do_constant_folding=True, # 优化常量 input_names=['input'], # 输入节点名 output_names=['loc', 'conf', 'landms'], # 输出节点名 dynamic_axes={'input': {0: 'batch_size'}, # 支持动态batch 'loc': {0: 'batch_size'}, 'conf': {0: 'batch_size'}, 'landms': {0: 'batch_size'}} ) print("模型已导出为 retinaface.onnx")转换成功后,你就可以使用ONNX Runtime等工具在各种平台上高效地运行这个人脸检测模型了。
5. 写在最后
走完这一整套流程,从准备数据、处理标注、增强图像,到配置损失函数、监控训练,最后部署模型,你应该对人脸检测模型的开发有了一个比较完整的体验。RetinaFace作为一个经典且强大的模型,把它训练好本身就是一个很有价值的项目。
实际做的时候,可能会遇到各种问题,比如数据标注不对齐、训练损失不收敛、推理速度慢等等。这都是正常的,解决问题的过程就是学习的过程。我建议你先在一个小规模的数据子集上跑通整个流程,确保代码没问题,再扩展到全量数据上训练,这样能节省很多时间。
训练模型有点像养花,需要耐心调整各种“养分”(参数)。多观察损失曲线,多分析模型在验证集上的错误案例(比如哪些小人脸没检测到,哪些框不准),然后有针对性地调整数据增强策略、损失权重或者模型结构,你的模型就会越来越“聪明”。
获取更多AI镜像
想探索更多AI镜像和应用场景?访问 CSDN星图镜像广场,提供丰富的预置镜像,覆盖大模型推理、图像生成、视频生成、模型微调等多个领域,支持一键部署。