☰
知识蒸馏+人脸关键点检测:本科毕设极小模型方案
2026/10/5 8:12:23 网站建设 项目流程

简介:这份资源是面向计算机、人工智能、通信工程等专业在校学生与教师的本科毕业设计项目源码,核心目标是通过知识蒸馏训练一个用于人脸关键点检测的极小模型。项目将大模型的知识迁移到轻量网络中,在保证精度的同时压缩模型体积,适合作为毕设、课程设计或项目初期立项演示,也适合具备一定Python基础的学习者进阶研究。压缩包共约2000个文件,整体约408.9MB,其中997个png图像与987个pts标注文件构成训练与测试数据集,11个py脚本承载模型定义、蒸馏训练与推理流程,另有csv、json配置及md说明文档,目录结构清晰,便于按模块阅读与复现。目前已有76人学习关注。代码均经过实际运行验证,读者可据此掌握知识蒸馏的完整实现思路、数据组织方式与关键点检测的评估方法,并在此基础上修改扩展,用于自身课题或功能迭代。

1. 知识蒸馏 + 人脸关键点检测:一个本科毕设能落地的极小模型方案

人脸关键点检测要在手机端或嵌入式设备上跑,模型体积和推理速度是硬约束。直接拿 ResNet 或 HRNet 当 backbone,精度是够了,但参数量动辄几十兆,部署到端侧基本不现实。本科毕业设计如果选这个方向,核心矛盾就是:怎么在有限算力下,让一个小模型学到接近大模型的检测能力。知识蒸馏就是解决这个矛盾的常见手段——用一个已经训练好的大模型(教师)去指导小模型(学生)训练,让学生在不增加推理成本的前提下逼近教师的输出分布。这个方案适合计算机视觉方向的本科毕设,需要 Python、PyTorch 基础,做完能拿到一套可复现的训练流程和可展示的对比实验。

2. 教师-学生架构怎么搭:从选型到最小可跑代码

2.1 为什么人脸关键点检测适合做蒸馏

人脸关键点检测的输出是固定数量的坐标点,常见的是 5 点(双眼、鼻尖、嘴角)或 68 点(dlib 标准)。这个任务有两个特点让它特别适合知识蒸馏:第一,输出是结构化的坐标回归,教师模型输出的软标签(soft label)比硬标签(ground truth 坐标)携带更多信息,比如某个点在遮挡情况下的概率分布;第二,教师和学生可以共享同样的输入分辨率,蒸馏时不需要额外的数据增强对齐。

常见做法是教师用 HRNet-W18 或 ResNet-50 加一个回归头,学生用 MobileNetV3-Small 或自己搭一个 4 层卷积加全连接的小网络。教师参数量大概 10M~30M,学生控制在 0.5M~2M 之间。蒸馏的目标函数一般由两部分组成:学生预测和 ground truth 的 MSE 损失,加上学生和教师输出之间的蒸馏损失。蒸馏损失可以用 KL 散度(对输出做 softmax 后)或直接 MSE(对坐标回归更直接)。

我一般会先用教师模型在数据集上训到收敛,保存权重,然后冻结教师,只训学生。这样避免教师和学生同时更新导致的训练不稳定。

2.2 教师模型的最小实现

教师模型不需要太复杂,本科毕设层面用 ResNet-18 改一个回归头就够了。下面是一个可直接跑的教师网络定义:

import torch import torch.nn as nn import torchvision.models as models class TeacherNet(nn.Module): def __init__(self, num_keypoints=68): super().__init__() # 用预训练 ResNet-18 做 backbone self.backbone = models.resnet18(weights=models.ResNet18_Weights.DEFAULT) # 替换最后的全连接层,输出 num_keypoints * 2(x,y 坐标) in_features = self.backbone.fc.in_features self.backbone.fc = nn.Linear(in_features, num_keypoints * 2) def forward(self, x): # 输出 shape: (batch, num_keypoints * 2) return self.backbone(x)

逻辑说明:backbone 用 torchvision 的预训练 ResNet-18,去掉原来的 1000 类分类头,换成一个线性层输出 136 维(68 点 × 2)。输入图像统一 resize 到 224×224,归一化用 ImageNet 的均值和方差。参数上,num_keypoints 根据数据集定,300W-LP 是 68 点,AFLW 是 21 点,自己标的数据集按实际点数改。

训练教师时用 MSE 损失,优化器选 Adam,学习率 1e-3,batch size 32,跑 100 个 epoch 左右。如果显存不够,把 batch size 降到 16,学习率对应降到 5e-4。

2.3 学生模型和蒸馏损失

学生模型要足够小,但也不能小到学不动。我一般用 MobileNetV3-Small 的 backbone,后面接两层全连接:

import torch import torch.nn as nn import torchvision.models as models class StudentNet(nn.Module): def __init__(self, num_keypoints=68): super().__init__() self.backbone = models.mobilenet_v3_small(weights=models.MobileNet_V3_Small_Weights.DEFAULT) # MobileNetV3-Small 最后一层特征维度是 576 self.backbone.classifier = nn.Sequential( nn.Linear(576, 256), nn.ReLU(), nn.Linear(256, num_keypoints * 2) ) def forward(self, x): return self.backbone(x)

蒸馏训练的核心代码:

import torch import torch.nn as nn import torch.nn.functional as F def distillation_loss(student_out, teacher_out, gt, alpha=0.5, temperature=2.0): """ student_out: 学生输出 (batch, 136) teacher_out: 教师输出 (batch, 136) gt: ground truth 坐标 (batch, 136) alpha: 蒸馏损失权重 temperature: 温度系数,用于软化教师输出 """ # 硬标签损失:学生 vs ground truth hard_loss = F.mse_loss(student_out, gt) # 软标签损失:学生 vs 教师(温度缩放) soft_student = student_out / temperature soft_teacher = teacher_out / temperature soft_loss = F.mse_loss(soft_student, soft_teacher) # 加权组合 return alpha * soft_loss + (1 - alpha) * hard_loss

逻辑说明:hard_loss 让学生直接拟合真实坐标,soft_loss 让学生模仿教师的输出分布。temperature 控制软化的程度,坐标回归任务里 temperature 设 1.0~3.0 都行,我一般用 2.0。alpha 是蒸馏损失的权重,设 0.5 意味着两者同等重要;如果教师训得特别好,可以把 alpha 提到 0.7;如果教师本身精度一般,降到 0.3 避免学生被带偏。

训练循环里,每个 batch 先跑教师前向(不需要梯度),再跑学生前向,然后计算蒸馏损失反传。教师权重全程冻结,用teacher.eval()和torch.no_grad()包住教师的前向。

3. 数据准备和训练流程:从 300W 到自定义数据集

3.1 数据集选择和标注格式转换

人脸关键点检测常用的公开数据集有 300W、AFLW、COFW。300W 包含 68 点标注,覆盖不同姿态、表情、光照,适合做毕设。下载后标注是 .pts 格式,每行一个点的 x y 坐标。需要转成 PyTorch 能读的格式,我一般转成 CSV 或直接存成 numpy 数组。

import numpy as np import os def parse_pts(pts_path): """解析 .pts 文件,返回 (68, 2) 的 numpy 数组""" with open(pts_path, 'r') as f: lines = f.readlines() # 跳过前 3 行头部信息 coords = [] for line in lines[3:]: line = line.strip() if line.startswith('}') or not line: break x, y = line.split() coords.append([float(x), float(y)]) return np.array(coords) def convert_dataset(root_dir, output_csv): """遍历数据集目录,把所有样本转成一个 CSV""" records = [] for subdir in os.listdir(root_dir): img_dir = os.path.join(root_dir, subdir) if not os.path.isdir(img_dir): continue for fname in os.listdir(img_dir): if fname.endswith('.pts'): pts_path = os.path.join(img_dir, fname) img_path = pts_path.replace('.pts', '.jpg') if not os.path.exists(img_path): continue coords = parse_pts(pts_path) # 展平成 136 维 flat = coords.flatten().tolist() records.append([img_path] + flat) # 写 CSV import csv with open(output_csv, 'w', newline='') as f: writer = csv.writer(f) header = ['image_path'] + [f'p{i}' for i in range(136)] writer.writerow(header) writer.writerows(records) print(f'转换完成,共 {len(records)} 个样本')

逻辑说明:parse_pts 跳过前 3 行头部,读到}结束。convert_dataset 遍历所有子目录,把 .pts 和对应 .jpg 配对,坐标展平成 136 维存 CSV。参数上,root_dir 是 300W 解压后的根目录,output_csv 指定输出路径。注意有些样本可能没有对应图片,代码里做了存在性检查。

3.2 数据增强和 DataLoader

人脸关键点检测对几何变换敏感,增强时坐标要同步变换。常见做法是随机水平翻转、随机旋转 ±15 度、随机缩放 0.9~1.1 倍。水平翻转时要注意左右点的索引要对调,68 点的对调关系是固定的。

import torch from torch.utils.data import Dataset, DataLoader import cv2 import numpy as np class FaceKeypointDataset(Dataset): def __init__(self, csv_path, img_size=224, augment=True): import pandas as pd self.df = pd.read_csv(csv_path) self.img_size = img_size self.augment = augment def __len__(self): return len(self.df) def __getitem__(self, idx): row = self.df.iloc[idx] img = cv2.imread(row['image_path']) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) coords = row.values[1:].astype(np.float32).reshape(-1, 2) # resize 到统一尺寸 h, w = img.shape[:2] img = cv2.resize(img, (self.img_size, self.img_size)) coords[:, 0] *= self.img_size / w coords[:, 1] *= self.img_size / h # 归一化到 [0,1] coords[:, 0] /= self.img_size coords[:, 1] /= self.img_size # 转 tensor img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0 coords = torch.from_numpy(coords.flatten()).float() return img, coords

逻辑说明:getitem读图、转 RGB、resize 到 224×224,坐标按比例缩放后归一化到 [0,1]。返回的 coords 是 136 维向量。DataLoader 的 batch_size 设 32,shuffle=True,num_workers 设 4(Windows 下设 0 避免多进程问题)。

3.3 完整训练循环

把教师、学生、数据加载和蒸馏损失串起来:

import torch from torch.utils.data import DataLoader from tqdm import tqdm def train_student(teacher, student, train_loader, val_loader, epochs=100, lr=1e-3, alpha=0.5, temperature=2.0): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') teacher = teacher.to(device).eval() student = student.to(device) optimizer = torch.optim.Adam(student.parameters(), lr=lr) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs) for epoch in range(epochs): student.train() total_loss = 0 for imgs, gts in tqdm(train_loader, desc=f'Epoch {epoch+1}'): imgs, gts = imgs.to(device), gts.to(device) with torch.no_grad(): teacher_out = teacher(imgs) student_out = student(imgs) loss = distillation_loss(student_out, teacher_out, gts, alpha, temperature) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() scheduler.step() avg_loss = total_loss / len(train_loader) print(f'Epoch {epoch+1}, Loss: {avg_loss:.6f}') # 每 10 个 epoch 验证一次 if (epoch + 1) % 10 == 0: student.eval() val_loss = 0 with torch.no_grad(): for imgs, gts in val_loader: imgs, gts = imgs.to(device), gts.to(device) out = student(imgs) val_loss += torch.nn.functional.mse_loss(out, gts).item() print(f'Val Loss: {val_loss / len(val_loader):.6f}') return student

逻辑说明:教师用 eval 模式且 no_grad 包住前向,不更新参数。学生用 train 模式,蒸馏损失反传。学习率用余弦退火,从 1e-3 降到接近 0。每 10 个 epoch 在验证集上算 MSE,监控是否过拟合。参数上,epochs 设 100 起步,如果 loss 还在降可以加到 150;alpha 和 temperature 按 2.3 节的建议调。

4. 避坑与排查:蒸馏训练里最容易翻车的 5 个点

4.1 教师输出没归一化,学生学了个寂寞

现象:蒸馏 loss 一直不降,学生输出和教师输出差距很大。原因:教师训练时坐标归一化到 [0,1],但学生训练时 ground truth 没做同样归一化,导致 soft_loss 和 hard_loss 量级差几个数量级。解决:确保教师和学生的输入预处理、坐标归一化方式完全一致。我一般把归一化写在 Dataset 里,教师和学生共用同一个 DataLoader。

4.2 温度系数设太高,坐标回归变成分类

现象:学生预测的坐标全部挤在图像中心附近,方差很小。原因:temperature 设得太大(比如 10),softmax 后教师输出变得接近均匀分布,学生学不到有效的坐标信息。解决:坐标回归任务里 temperature 不要超过 3.0,我一般用 1.0~2.0。如果发现学生输出方差太小,先把 temperature 降到 1.0 试试。

4.3 教师太强或太弱,蒸馏反而拖后腿

现象:加了蒸馏 loss 后,学生精度比不加还差。原因:教师本身在验证集上精度就不高(比如 MSE 大于 0.01),或者教师和学生容量差距太大(教师 30M,学生 0.1M),学生根本模仿不了。解决:先确认教师收敛且验证集 MSE 在合理范围(68 点任务一般小于 0.005)。如果教师没问题但学生学不动,把 alpha 降到 0.3,或者换一个稍大的学生模型(比如 MobileNetV2)。

4.4 数据增强时坐标没同步变换

现象:训练 loss 震荡严重,验证集精度远低于训练集。原因:图像做了随机旋转或翻转,但坐标没跟着变,导致标签和图像不对应。解决:所有几何变换必须同时作用于图像和坐标。水平翻转时,68 点的左右对调索引是固定的,比如点 0 和点 16 对调,点 1 和点 15 对调,以此类推。建议先用小批量数据可视化检查增强后的图像和坐标是否对齐。

4.5 显存不够导致 batch size 太小,训练不稳定

现象:CUDA out of memory,只能把 batch size 降到 4 或 8,loss 波动大。原因:教师和学生同时加载到显存,加上中间激活值,显存占用翻倍。解决:教师前向用torch.no_grad()且不保存中间激活;如果还不够,把教师推理放在 CPU 上,只把输出传到 GPU。或者用梯度累积,每 4 个 batch 更新一次参数,等效 batch size 32。

5. 蒸馏效果验证与模型导出:怎么证明学生真的学到了

训练完之后,需要一套验证流程来证明蒸馏确实有效。我一般做三组对比:学生不加蒸馏直接训、学生加蒸馏训、教师直接训。在同一个验证集上算 MSE 和 NME(归一化平均误差),NME 的计算方式是预测点和真实点的欧氏距离除以两眼间距或人脸框对角线长度。

import torch import numpy as np def compute_nme(pred, gt, norm_factor): """ pred: (N, 68, 2) 预测坐标 gt: (N, 68, 2) 真实坐标 norm_factor: (N,) 每个样本的归一化因子,通常用两眼间距 """ pred = pred.reshape(-1, 68, 2) gt = gt.reshape(-1, 68, 2) diff = pred - gt dist = np.sqrt((diff ** 2).sum(axis=2)) # (N, 68) nme = dist.mean(axis=1) / norm_factor # (N,) return nme.mean() # 假设已经跑完验证集,拿到 preds 和 gts # norm_factor 用两眼间距:点 36 到点 45 的距离(68 点标准)

逻辑说明:NME 是行业常用指标,比 MSE 更直观。norm_factor 用两眼间距或人脸框对角线,68 点标准里点 36 和点 45 是左右眼外眼角。算出来的 NME 一般在 0.05~0.08 之间算不错,低于 0.05 算很好。

模型导出用 ONNX:

import torch.onnx student.eval() dummy_input = torch.randn(1, 3, 224, 224).to(device) torch.onnx.export( student, dummy_input, 'student_keypoint.onnx', input_names=['input'], output_names=['output'], dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}}, opset_version=11 )

导出后可以用 onnxruntime 验证推理结果和 PyTorch 一致。学生模型导出后大概 2~5MB,推理一张 224×224 的图在 CPU 上 10~20ms,满足端侧部署要求。

最后说一个我踩过的坑:蒸馏训练时不要过早停止。教师的知识迁移需要时间,前 30 个 epoch 学生可能看起来没学到什么,loss 下降很慢,但 50 epoch 之后会突然变好。我一般至少跑 100 epoch,中间不要因为验证 loss 波动就中断。希望帮到你。

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

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

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

立即咨询