基于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张人脸图像。下载后需要按以下步骤预处理:

  1. 目录结构调整 :确保每个身份对应一个独立文件夹
  2. 人脸检测与对齐 :使用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 训练循环实现

完整的训练流程包含以下关键步骤:

  1. 学习率调度 :使用余弦退火策略
  2. 模型保存 :基于验证集准确率保存最佳模型
  3. 混合精度训练 :加速训练过程
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%。在实际项目中,这种平衡往往是可以接受的。

Logo

汇聚全球AI编程工具,助力开发者即刻编程。

更多推荐