用MobileNetV1复现FaceNet人脸识别:从PyTorch代码到CASIA-WebFace数据集实战
·
基于MobileNetV1的FaceNet人脸识别实战:从零构建PyTorch模型到CASIA-WebFace训练全流程
人脸识别技术正逐渐渗透到日常生活的各个角落,从手机解锁到门禁系统,其核心在于如何将人脸图像转化为具有判别性的特征向量。本文将带你完整实现一个基于MobileNetV1的FaceNet模型,使用PyTorch框架和CASIA-WebFace数据集,从环境搭建到模型训练,逐步解析每个关键环节。
1. 环境准备与数据预处理
1.1 基础环境配置
首先需要准备Python 3.8+环境和必要的库。推荐使用conda创建虚拟环境:
conda create -n facenet python=3.8
conda activate facenet
pip install torch torchvision torchaudio
pip install opencv-python pillow matplotlib tqdm
对于GPU加速,建议安装对应CUDA版本的PyTorch。可以通过以下命令检查GPU是否可用:
import torch
print(torch.cuda.is_available()) # 应输出True
print(torch.__version__) # 建议1.8.0+
1.2 CASIA-WebFace数据集处理
CASIA-WebFace包含10,575个身份的494,414张人脸图像。下载后需要按以下步骤预处理:
- 目录结构调整 :确保每个身份对应一个独立文件夹
- 人脸检测与对齐 :使用MTCNN进行人脸检测和关键点对齐
from facenet_pytorch import MTCNN
import cv2
import os
mtcnn = MTCNN(keep_all=True)
def align_face(image_path, save_path):
img = cv2.imread(image_path)
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
boxes, probs = mtcnn.detect(img_rgb)
if boxes is not None:
box = boxes[0]
cv2.rectangle(img, (int(box[0]), int(box[1])),
(int(box[2]), int(box[3])), (0,255,0), 2)
cv2.imwrite(save_path, img)
注意:对齐后的人脸图像建议统一调整为160×160像素,这是MobileNetV1的标准输入尺寸
2. MobileNetV1骨干网络实现
2.1 深度可分离卷积实现
MobileNetV1的核心是深度可分离卷积,它将标准卷积分解为深度卷积和点卷积:
import torch.nn as nn
class DepthwiseSeparableConv(nn.Module):
def __init__(self, in_channels, out_channels, stride):
super().__init__()
self.depthwise = nn.Sequential(
nn.Conv2d(in_channels, in_channels, 3, stride, 1, groups=in_channels, bias=False),
nn.BatchNorm2d(in_channels),
nn.ReLU6(inplace=True)
)
self.pointwise = nn.Sequential(
nn.Conv2d(in_channels, out_channels, 1, 1, 0, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU6(inplace=True)
)
def forward(self, x):
x = self.depthwise(x)
x = self.pointwise(x)
return x
2.2 完整MobileNetV1架构
基于上述模块构建完整的MobileNetV1:
class MobileNetV1(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 32, 3, 2, 1, bias=False),
nn.BatchNorm2d(32),
nn.ReLU6(inplace=True),
DepthwiseSeparableConv(32, 64, 1),
DepthwiseSeparableConv(64, 128, 2),
DepthwiseSeparableConv(128, 128, 1),
DepthwiseSeparableConv(128, 256, 2),
DepthwiseSeparableConv(256, 256, 1),
DepthwiseSeparableConv(256, 512, 2),
# 重复6次512通道的深度可分离卷积
*[DepthwiseSeparableConv(512, 512, 1) for _ in range(6)],
DepthwiseSeparableConv(512, 1024, 2),
DepthwiseSeparableConv(1024, 1024, 1)
)
self.avg_pool = nn.AdaptiveAvgPool2d((1,1))
def forward(self, x):
x = self.features(x)
x = self.avg_pool(x)
return x.view(x.size(0), -1)
提示:MobileNetV1相比原始FaceNet使用的Inception-ResNet体积缩小约10倍,适合移动端部署
3. FaceNet模型构建与损失函数
3.1 特征嵌入层实现
在MobileNetV1基础上添加嵌入层和L2标准化:
class FaceNet(nn.Module):
def __init__(self, embedding_size=128, num_classes=None):
super().__init__()
self.backbone = MobileNetV1()
self.bottleneck = nn.Linear(1024, embedding_size, bias=False)
self.bn = nn.BatchNorm1d(embedding_size, eps=0.001, momentum=0.1)
if num_classes is not None:
self.classifier = nn.Linear(embedding_size, num_classes)
def forward(self, x):
x = self.backbone(x)
x = self.bottleneck(x)
x = self.bn(x)
x = nn.functional.normalize(x, p=2, dim=1)
return x
3.2 Triplet Loss与交叉熵组合
实现三元组损失并配合交叉熵辅助训练:
class TripletLoss(nn.Module):
def __init__(self, margin=0.2):
super().__init__()
self.margin = margin
def forward(self, anchor, positive, negative):
pos_dist = torch.sum((anchor - positive)**2, dim=1)
neg_dist = torch.sum((anchor - negative)**2, dim=1)
losses = torch.relu(pos_dist - neg_dist + self.margin)
return losses.mean()
class CombinedLoss(nn.Module):
def __init__(self, triplet_weight=1.0, ce_weight=0.5):
super().__init__()
self.triplet = TripletLoss()
self.ce = nn.CrossEntropyLoss()
self.triplet_weight = triplet_weight
self.ce_weight = ce_weight
def forward(self, anchor, positive, negative, logits, labels):
triplet_loss = self.triplet(anchor, positive, negative)
ce_loss = self.ce(logits, labels)
return self.triplet_weight*triplet_loss + self.ce_weight*ce_loss
4. 训练流程与调优技巧
4.1 数据加载与三元组采样
高效的三元组采样对训练至关重要:
from torch.utils.data import Dataset
import random
class TripletDataset(Dataset):
def __init__(self, root_dir, transform=None):
self.classes = os.listdir(root_dir)
self.class_to_idx = {c:i for i,c in enumerate(self.classes)}
self.samples = []
for c in self.classes:
class_dir = os.path.join(root_dir, c)
for img_name in os.listdir(class_dir):
self.samples.append((os.path.join(class_dir, img_name), self.class_to_idx[c]))
def __getitem__(self, index):
anchor_path, anchor_label = self.samples[index]
# 随机选择正样本
positive_samples = [s for s in self.samples if s[1]==anchor_label]
positive_path, _ = random.choice(positive_samples)
# 随机选择负样本
negative_samples = [s for s in self.samples if s[1]!=anchor_label]
negative_path, _ = random.choice(negative_samples)
anchor = load_and_transform(anchor_path)
positive = load_and_transform(positive_path)
negative = load_and_transform(negative_path)
return anchor, positive, negative, anchor_label
def __len__(self):
return len(self.samples)
4.2 训练循环实现
完整的训练流程包含以下关键步骤:
- 学习率调度 :使用余弦退火策略
- 模型保存 :基于验证集准确率保存最佳模型
- 混合精度训练 :加速训练过程
from torch.optim import Adam
from torch.optim.lr_scheduler import CosineAnnealingLR
def train_model(model, train_loader, val_loader, epochs=50):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
optimizer = Adam(model.parameters(), lr=0.001)
scheduler = CosineAnnealingLR(optimizer, T_max=epochs)
criterion = CombinedLoss()
best_acc = 0.0
for epoch in range(epochs):
model.train()
for batch in train_loader:
anchor, pos, neg, labels = [x.to(device) for x in batch]
optimizer.zero_grad()
anchor_emb = model(anchor)
pos_emb = model(pos)
neg_emb = model(neg)
logits = model.classifier(anchor_emb)
loss = criterion(anchor_emb, pos_emb, neg_emb, logits, labels)
loss.backward()
optimizer.step()
scheduler.step()
# 验证集评估
val_acc = evaluate(model, val_loader)
if val_acc > best_acc:
best_acc = val_acc
torch.save(model.state_dict(), 'best_model.pth')
print(f'Epoch {epoch+1}/{epochs} | Loss: {loss.item():.4f} | Val Acc: {val_acc:.2f}%')
4.3 常见问题与解决方案
在实际训练中可能会遇到以下典型问题:
| 问题现象 | 可能原因 | 解决方案 |
|---|---|---|
| 损失不下降 | 学习率过大/过小 | 尝试0.0001-0.01范围内的学习率 |
| 准确率波动大 | 批次内三元组质量差 | 增加批次大小或使用在线困难样本挖掘 |
| 过拟合明显 | 模型容量过大 | 增加Dropout或减少嵌入维度 |
| 训练速度慢 | 未使用GPU加速 | 检查CUDA配置或减小输入尺寸 |
提示:对于小型数据集,可以冻结MobileNetV1的部分层,只训练顶层参数
5. 模型评估与应用部署
5.1 人脸验证评估
在LFW数据集上评估模型性能:
from sklearn.metrics import roc_curve, auc
def evaluate_verification(model, pairs_path):
model.eval()
distances, labels = [], []
with open(pairs_path) as f:
for line in f:
img1, img2, same = line.strip().split()
emb1 = model(load_image(img1).unsqueeze(0))
emb2 = model(load_image(img2).unsqueeze(0))
dist = torch.norm(emb1 - emb2, p=2).item()
distances.append(dist)
labels.append(int(same))
fpr, tpr, _ = roc_curve(labels, distances)
roc_auc = auc(fpr, tpr)
return roc_auc
5.2 实时人脸识别部署
使用OpenCV实现实时人脸识别:
import cv2
from PIL import Image
def live_recognition(model, threshold=0.7):
cap = cv2.VideoCapture(0)
mtcnn = MTCNN()
known_faces = {} # 预存已知人脸特征
while True:
ret, frame = cap.read()
if not ret: break
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
boxes, _ = mtcnn.detect(frame_rgb)
if boxes is not None:
for box in boxes:
x1, y1, x2, y2 = map(int, box)
face = frame_rgb[y1:y2, x1:x2]
face = Image.fromarray(face).resize((160,160))
emb = model(transform(face).unsqueeze(0))
# 计算与已知人脸的相似度
match = False
for name, known_emb in known_faces.items():
dist = torch.norm(emb - known_emb, p=2).item()
if dist < threshold:
cv2.putText(frame, name, (x1,y1-10),
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0,255,0), 2)
match = True
break
if not match:
cv2.putText(frame, 'Unknown', (x1,y1-10),
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0,0,255), 2)
cv2.rectangle(frame, (x1,y1), (x2,y2), (0,255,0), 2)
cv2.imshow('Face Recognition', frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
5.3 模型量化与优化
为移动端部署进行模型量化:
model = FaceNet().eval()
quantized_model = torch.quantization.quantize_dynamic(
model, {nn.Linear}, dtype=torch.qint8
)
torch.jit.save(torch.jit.script(quantized_model), 'facenet_quantized.pt')
经过量化后,模型大小可缩减至原来的1/4左右,推理速度提升2-3倍,而准确率损失通常不超过2%。在实际项目中,这种平衡往往是可以接受的。
更多推荐

所有评论(0)