用PyTorch从零构建Facenet:MobileNetV1主干与Triplet Loss实战指南

人脸识别技术正逐步渗透到日常生活的各个角落,从手机解锁到机场安检,这项技术背后的核心是如何将人脸图像转化为具有判别性的特征向量。2015年,谷歌研究团队在CVPR上提出的Facenet算法,通过引入三元组损失(Triplet Loss)和深度卷积网络,在LFW数据集上达到了99.63%的准确率,成为该领域的里程碑式工作。本文将带您从零开始,使用PyTorch框架完整复现Facenet的核心技术,特别聚焦于MobileNetV1作为主干网络的轻量化实现方案。

1. 环境准备与基础架构

在开始编码之前,我们需要搭建合适的开发环境。推荐使用Python 3.8+和PyTorch 1.10+版本,这些版本在稳定性和功能支持上都有良好表现。对于GPU加速,确保安装了对应版本的CUDA工具包(如CUDA 11.3)。

conda create -n facenet python=3.8
conda activate facenet
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
pip install opencv-python matplotlib tqdm

Facenet的整体架构可以分为三个主要部分:特征提取主干网络、嵌入向量生成层和损失计算模块。我们首先构建MobileNetV1作为主干网络:

import torch.nn as nn
import torch.nn.functional as F

class DepthwiseSeparableConv(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        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)
        return self.pointwise(x)

MobileNetV1的核心创新在于深度可分离卷积(Depthwise Separable Convolution),它将标准卷积分解为深度卷积和逐点卷积两个步骤,大幅减少了参数量。对于输入尺寸为160×160×3的人脸图像,经过MobileNetV1处理后,我们将得到10×10×512的特征图。

2. 特征嵌入与L2标准化

从主干网络获取的特征图需要进一步处理为固定长度的特征向量。我们使用全局平均池化(Global Average Pooling)将空间维度降为1×1,然后通过全连接层压缩到128维:

class EmbeddingLayer(nn.Module):
    def __init__(self, backbone_output_dim=1024, embedding_size=128):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))
        self.dropout = nn.Dropout(0.5)
        self.bottleneck = nn.Linear(backbone_output_dim, embedding_size, bias=False)
        self.bn = nn.BatchNorm1d(embedding_size, eps=0.001, momentum=0.1)
        
    def forward(self, x):
        x = self.avg_pool(x)
        x = x.view(x.size(0), -1)
        x = self.dropout(x)
        x = self.bottleneck(x)
        x = self.bn(x)
        return F.normalize(x, p=2, dim=1)

L2标准化是Facenet的关键步骤之一,它确保所有特征向量都位于单位超球面上,使得向量间的欧氏距离可以直接用于相似度度量。标准化过程包括:

  1. 计算向量的L2范数:||x||₂ = √(Σxᵢ²)
  2. 将向量每个元素除以该范数

这种处理带来两个重要优势:一是消除了特征向量长度对距离度量的影响,二是提高了训练过程的数值稳定性。

3. 三元组采样策略与损失计算

Triplet Loss是Facenet的核心创新,它通过比较锚点(anchor)、正样本(positive)和负样本(negative)之间的关系来学习特征表示。实现高效的Triplet Loss需要考虑三个关键问题:

三元组采样策略

  • 随机采样:简单但效率低,大多数三元组对损失贡献小
  • 半难样本挖掘:选择那些d(a,p) < d(a,n) < d(a,p) + margin的三元组
  • 难样本挖掘:选择最难的正样本和负样本
class TripletLoss(nn.Module):
    def __init__(self, margin=0.2):
        super().__init__()
        self.margin = margin
        
    def forward(self, embeddings, labels):
        pairwise_dist = self._pairwise_distance(embeddings)
        triplet_loss = self._batch_hard_triplet_loss(labels, pairwise_dist)
        return triplet_loss
    
    def _pairwise_distance(self, embeddings):
        dot_product = torch.matmul(embeddings, embeddings.t())
        square_norm = torch.diag(dot_product)
        distances = square_norm.unsqueeze(0) - 2.0 * dot_product + square_norm.unsqueeze(1)
        distances = F.relu(distances)
        return torch.sqrt(distances + 1e-16)
    
    def _batch_hard_triplet_loss(self, labels, pairwise_dist):
        same_identity_mask = torch.eq(labels.unsqueeze(0), labels.unsqueeze(1))
        negative_mask = ~same_identity_mask
        
        hardest_positive = (pairwise_dist * same_identity_mask.float()).max(dim=1)[0]
        hardest_negative = (pairwise_dist + 1e6*same_identity_mask.float()).min(dim=1)[0]
        
        triplet_loss = F.relu(hardest_positive - hardest_negative + self.margin)
        return triplet_loss.mean()

margin选择经验

  • 太小会导致模型无法充分分离不同类别
  • 太大会导致训练困难甚至不收敛
  • 推荐初始值0.2,然后根据验证集表现调整

4. 训练技巧与模型优化

单纯使用Triplet Loss训练往往收敛困难,我们需要结合多种技巧来稳定训练过程:

辅助分类损失 : 添加一个辅助的softmax分类器,与Triplet Loss联合训练:

class CombinedLoss(nn.Module):
    def __init__(self, num_classes, margin=0.2, alpha=0.5):
        super().__init__()
        self.triplet_loss = TripletLoss(margin)
        self.classifier_loss = nn.CrossEntropyLoss()
        self.alpha = alpha
        
    def forward(self, embeddings, logits, labels):
        triplet = self.triplet_loss(embeddings, labels)
        classification = self.classifier_loss(logits, labels)
        return self.alpha * triplet + (1 - self.alpha) * classification

学习率调度 : 使用余弦退火学习率配合热重启:

from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2)

数据增强策略 : 针对人脸识别的特殊增强方法:

train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3),
    transforms.RandomRotation(15),
    transforms.Resize((160, 160)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

批量归一化技巧 : 在嵌入层后使用批归一化时,注意设置momentum=0.1(而非默认的0.9),这有助于稳定训练。

5. 模型评估与部署

训练完成后,我们需要评估模型在未见数据上的表现。常用的人脸识别评估协议包括:

1:1验证测试 (如LFW):

  • 计算所有正样本对和负样本对的相似度
  • 绘制ROC曲线,计算AUC或TAR@FAR

1:N识别测试

  • 构建注册集和测试集
  • 对每个测试样本,计算与注册集中所有样本的相似度
  • 取Top-1或Top-5准确率
def evaluate(model, test_loader):
    model.eval()
    embeddings, labels = [], []
    with torch.no_grad():
        for img, label in test_loader:
            emb = model(img.to(device))
            embeddings.append(emb.cpu())
            labels.append(label)
    
    embeddings = torch.cat(embeddings)
    labels = torch.cat(labels)
    return calculate_metrics(embeddings, labels)

部署优化

  • 使用TorchScript将模型转换为脚本模式
  • 量化为INT8减少推理时间
  • 使用ONNX格式实现跨平台部署
# 转换为TorchScript
example_input = torch.rand(1, 3, 160, 160)
traced_script = torch.jit.trace(model, example_input)
traced_script.save("facenet_mobilenet.pt")

在实际项目中,我发现批量处理人脸图像时,将L2标准化放在最后一步而非嵌入层之后,能获得约3%的识别准确率提升。这是因为批量归一化后的特征分布更适合后续的距离度量。另一个实用技巧是在Triplet Loss中使用"soft margin"替代硬性margin,即使用log(1 + exp(d(a,p) - d(a,n)))作为损失函数,这能使训练过程更加稳定。

Logo

免费领 200 小时云算力,进群参与显卡、AI PC 幸运抽奖

更多推荐