一 代码分析
import argparse #用于解析命令行参数，比如train/match, --data, --epochs等
import math #用于数学计算，比如ArcFace 里面的cos, sin, pi
import os #用于创建目录，处理系统路径等
import random #用于设置Python 随即种子，保证实验可复现
from pathlib import path #更方便处理文件路径，比如path(args.out)
from PIL import Image # 用于读取图片文件
from tqdm import tqdm # 用于显示训练进度条
import torch #
import torch.nn as nn# Pytorch 神经网络模块，比如Linear, Conv2d, BatchNorm
import torch.nn.functional as F #pytorch函数式接口，比如normalize, linear
from torch.utils.data import DataLoader #用于批量加载数据
from torchvision import datasets, transforms #datasetsyong
from torchvision.modesl import resnet18 #使用torchvision自带的ResNet18 作为backbone
def set_seed(seed: int = 42): #定义设置随即种子的函数，默认为seed为42
random.seed(seed) 设置python内置random的随机种子
torch.manul_seed(seed) 设置pytorch CPU随机种子
torch.cuda.manul_seed_all(seed) 设置所有GPU的随机种子
class EnbeddingNet(nn.Module): 定义人脸特征提取网络，继承Pytorch的nn.Module
## 传统CNN backbone ResNet18 改成适合112x112人脸输入
#输出embedding 用于人脸匹配
def __init__(self, emb_dim: int = 512): 初始化网络，emb_dim 表示最终人脸特征维度 = 512
super().__int__() #调用父类nn.Module的初始化函数
try: #尝试使用新版本torchvision 的写法
net = resnet18(weights=None)#创建一个不加载预训练权重的ResNet18
except TypeError: #如果当前torchvision 版本不支持weights 参数，就走旧版本写法
net = resnet18(pretrained=false)#旧版本torchvision中用
pretrained=False 表示不加载预测训练权重
#原始 ResNet18 的conv1是7x7 stride = 2, 适合ImageNet的224x224图片
#人脸常用112x112输入，所以这里改成3x3 stride = 1 减少早期信息损失
net.conv1 = nn.Conv2d(#替换ResNet18 的第一层卷积
3, #输入通道数
64, #输出通道数，保持ResNet18原始设定
kernel_size = 3, 3x3 卷积
stride=1, # 步长改成 1，避免过早下采样
padding=1, # padding=1 可以保持特征图尺寸不变
bias=False # 后面有 BatchNorm，所以卷积层不需要 bias
)
net.maxpool = nn.Indetity() #移除ResNet18原始maxpool 下采样更少，保留更多人脸细节
in_features = net.fc.in_features #获取ResNet18 最后一层全脸阶层的输入维度
net.fc = nn.Linear( #替换ResNet18原始分类层
in_features, #输入维度来自ResNet18 backbone
emb_dim, #输出维度是人脸embedding维度，比如512
bias = False #人脸特征层经常不使用bias
)
self.backbone = net #保存修改后的ResNet18 作为特征提取主干网络
self.bn = nn.BatchNorm1d(emb_dim) #对最终embedding做batchNorm,稳定训练
def forward(self, x) #定义前向传播，x是输入图片的batch
x = self.backbone(x) #图片经过ResNet18,得到emb_dim维特征
x = self.bn(x) #对embedding做batchNorm 归一化处理
return x #返回人脸特向量
class ArcMarginProduct(nn.Module): #定义ArcFace分类头
#ArcFace分类头
#训练时用身份分类监督，推理时不用这个head, 只用EmbeddingNet
def __init__( #初始化ArcFace分类头
self, #当前模块对象
in_features: int #输入特征维度，也就是embedding维度
out_features:int, 输出类别数，也就是训练集中的身份数量
s: float = 64.0 #ArcFace 的缩放因子scale，常用64
m: float = 0.5 #ArcFace 的角度margin, 常用0。5
easy_margin: bool = False,#是否使用easy margin版本
)：
super().__init__() #调用父类nn.Module初始化函数
self.in_features = in_features #保存输入特征维度
self.out_features = out_features #保存类别数量
self.s = s#保存缩放因子
self.m = n #保存margin
self.easy_margin = easy_margin #保存是否使用easy margin
self.weight = nn.Parameter( #定义分类权重矩阵，作为可训练参数
torch.empty(out_features, in_features) #权重形状为类别数x特征维度
)
nn.init.xavier_uniform_(self.weight)#使用Xavier均匀分布初始化分类权重
self.cos_m = math.cos(m) #预计算cos(m) 后面用于cos(theta + m)
self.sin_m = math.sin(m) #预计算sin(m)后面用于cos(theta +m)
self.th = math.cos(math.pi - m)#阈值，用于处理角度边界
self.mm = math.sin(math.pi - m) * m #边界修正项，避免数值不稳定
def forward(self, embeddings, labels): #ArcFace前向传播，输入embedding和真实身份标签
embeddings = F.normalize(embeddings, dim= 1) #对每个人脸特征做L2归一化
weights = F.normalize(self.weight, dim=1)#对每个类别的分类权重做L2归一化
cosin = F.linear(embeddings, weights)#计算embedding和每个类别权重之间的余弦相似度
sine = torch.sqrt((1.0 - cosine.pow(2)).clamp(0,1)#根据sin^2 + cos^2 = 1) 计算sine
phi = cosine * self.cos_m - sine * self.sin_m #根据cos(theta + m) 公式加入角度margin
if self.easy_margin: #如果使用easy margin
phi = torch.where(#根据条件选择phi cosine
cosine > 0, #当cosine 大于0时
phi, #使用加margin 后的结果
cosine #否则保持原始cosine
)
else: #如果不是用easy margin, 这是ArcFace常见默认写法
phi = torch.where(#根据阈值选择phi或修正后的cosine
cosine > self.th,#当cosine大于阈值时
phi, #使用margin后的结果
cosine - self.mm #否则使用边界修正，保证训练稳定
)
one_hot = torch.zeros_like(cosine)#创建和cosine同形状的全0矩阵
one_hot.scatter_(#将真实类别位置设置为1，
1，#在类别维度上scatter
labels.view(-1, 1) #将标签变成batch_size x1的形状
1.0 真实类别位置填1
)
logits = (one_hot *phi) + ((1.0 - one_hor)*cosine)#真实类别用phi,非真实类别用原始cosine
logits *= self.s #乘以缩放因子s, 让logits 数值范围适合交叉熵
return logits #返回ArcFace分类logits
def build_train_transform(img_size: int): #构建训练阶段的数据增强流程
return transforms.Compose([
transforms.Resize((img_size, img_size)), #将图片缩放到固定大小，比如112x112
transforms.RandomHorizontalFlip(p=0.5) #以50%概率随机水平翻转，增强左右脸鲁棒性
transforms.RandomAffine(#随机仿射变换，模拟轻微角度和位置变化
degress = 5, #随机旋转角度范围为+-5度
translate = (0.03, 0.03) #水平垂直方向最多平移3%
scale = (0.95, 1.05) #随机缩放到95%,到105%
),
transform.ColorJitter(#随机颜色都懂，增强光照鲁棒性
brightness = 0.15 #随机调度亮度
contrast = 0.15 #随机调整对比度
saturation = 0.10 #随机调整饱和度
hue = 0.03 #随机挑战色相
),
transforms.ToTensor(), 将PIL图片转换成Pytorch Tensor, 像素范围 0, 1
transforms.Normalize(#对图片做归一化
[0.5, 0.5, 0.5]#RGB三个通道的均值
[0.5, 0.5, 0.5] RGB三个通道的标准差
) #归一化像素大于变成 -1, 1
])
def build_eval_transform(img_size: int): # 构建验证/推理阶段的图片预处理流程
return transforms.Compose([ # 将多个图像变换组合起来
transforms.Resize((img_size, img_size)), # 将图片缩放到固定大小
transforms.ToTensor(), # 将 PIL 图片转成 Tensor
transforms.Normalize( # 使用和训练阶段一致的归一化方式
[0.5, 0.5, 0.5], # RGB 均值
[0.5, 0.5, 0.5] # RGB 标准差
), # 归一化到约 [-1, 1]
])
def train(args): #定义训练函数，args来自命令行参数
set_seed(args.seed)#设置随机种子，尽量保证结果可复现
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 如果有 GPU 就用 cuda，否则用 cpu
os.makedirs(args.out, exist_ok=True) # 创建模型保存目录，如果目录已存在也不报错 使用GPU
dataset = datasets.ImageFolder(#使用ImageFolder 读取训练数据
args.data, #训练数据根目录，要求每一个身份一个子文件夹
transform=build_train_transform(args.img_size),#对训练图片应用数据增强和预处理
)
if len(dataset.classes) < 2: # 如果训练集中身份类别少于 2 个
raise ValueError("训练集至少需要 2 个身份类别。") # 抛出错误，因为分类训练至少需要两个类别
loader = DataLoader(#创建数据加载器
dataset, #要加载的数据集
batch_size=args.batch_size, #每个batch的图片数量
shuffle=True,#每个epoch打乱数据顺序
num_workers=args.workers,#使用多少个子进程加载数据
pin_memory=torch.cuda.is_available(), #使用GPU时开启pin_memory 加速数据传输
drop_last=True,#丢弃组后一个不完整batahc,避免BatchNorm小batch不稳定.
)
num_classes = len(dataset.classes) #h欧去身份类别数量
model = EmbeddingNet(emb_dim=args.emb_dim).to(device)#创建人脸embedding网络并移动到CPU/GPU
head = ArcMarginProduct(#创建ArcFace分类头
in_features=args.emb_dim,#输入维度等于embedding维度
out_features=num_classes,#输出类别数等于训练集身份数
s=args.scale, #ArcFace缩放因子
m=args.margin #ArcFace Margin
).to(device)#将ArcFace分类头移动到CPU/GPU
optimizer = torch.optim.Adamw(#使用Adamw优化器
list(model.parameters()) + list(head.parameters()), #同时优化backbone,和ArcFace head
lr = args.lr #学习率
weight_decay = args.weight_decay, #权重衰减，用于正则化
)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(#使用余弦退火学习率调度器
optimizer, #要调度的优化器
T_max = args.epochs #一个完整余弦的周期的epoch数
)
scaler = torch.cuda.amp.GradScaler(#创建AMP梯度所缩放器，用于混合精度训练
enabled = (args.amp and device.type == 'cuda')#只有用户启用 --amp且使用GPU时才开启
)
ce_lass = nn.CrossEntropLOss() #定义交叉熵损失，用于身份分类训练
print(f"device: {device}") # 打印当前使用的设备
print(f"images: {len(dataset)}") # 打印训练图片总数
print(f"classes: {num_classes}") # 打印身份类别总数
print(f"save dir: {args.out}") # 打印模型保存路径
for epoch in range(1, arg.epochs + 1): #从第1个epoch训练到args.epochs
model.train() #将embedding 网络设置为训练模式
head.train() #将ArcFace head设置为训练模式
total_loss = 0.0 累计当前epoch的总loss
total_correct =0 #累计当前epoch分类正确的样本数
total_seen = 0 累计当前epoch已经处理样本数
pbar = tqdm(loader, desc=f"Epoch {epoch}/{args.epochs}") # 为当前 epoch 创建进度条
for imgs, labels in pbar: #遍历每个batch, imgs是图片，labels是身份标签
imgs = imgs.to(device, non_blocking=True) #将图片移动到CPU/GPU
labels = labels.to(device, non_blocking=True) # 将标签移动到 CPU/GPU
optimizer.zero_grad(set_to_none=True)#清空上一轮反向传播的梯度
with torch.cuda.amp.autocast(#自动混合精度上下文
enabled = (args.amp and device.type == 'cuda')#只有启用--amp且在GPU上才使用混合精度
)：
embeddings = model(imgs)#通过backbone提取人脸embedding
logits = head(embeddings, labels)#通过ArcFace head得到分类logits
loss = ce_loss(logits, labels)#计算交叉熵分类损失
scaler.scale(loss).backward() #对loss做缩放反向传播，避免FP16梯度下溢
scaler.step(optimizer) #使用缩放后的梯度更新参数
scaler.update() #更新GradScaler的缩放系数
with torch.no_grad() #下面只是统计准确率，不需要计算梯度
pred = logits.argmax(dim=1) #取每个样本logits最大的类别作为预测身份
corret = (pred == labels).sum().item() 统计当前batch预测正确数量
bs = imgs.size(0) 当前batch的样本数
total_loss += loss.item() * bs 了家总loss 乘以bs是为了按照样本数加权
total_correct += correct 雷佳预测正确样本数
total_seen += bs #累加已经处理样本数
pbar.set_postfix({ # 更新进度条后面的显示信息
"loss": f"{total_loss / total_seen:.4f}", # 当前 epoch 平均 loss
"acc": f"{total_correct / total_seen:.4f}", # 当前 epoch 平均分类准确率
"lr": f"{optimizer.param_groups[0]['lr']:.2e}", # 当前学习率
})
scheduler.step() #每个epoch结束更新学习率
scheduler.step() # 每个 epoch 结束后更新学习率
ckpt = { # 构造 checkpoint 字典，用于保存模型和必要信息
"backbone": model.state_dict(), # 保存 embedding 网络参数
"head": head.state_dict(), # 保存 ArcFace head 参数
"class_to_idx": dataset.class_to_idx, # 保存类别名到类别编号的映射
"classes": dataset.classes, # 保存类别名称列表
"emb_dim": args.emb_dim, # 保存 embedding 维度
"img_size": args.img_size, # 保存训练图片尺寸
}
torch.save(ckpt, Path(args.out) / "last.pt") #保存最新模型到last.pt
if epoch % args.save_every == 0 :#如果当前epoch到达保存间隔
torch.save( #额外保存一个带epoch变好的checkpoint
ckpt, #要保存的checkpoint
Path(args.out) / f"epoch_{epoch}.pt" # 保存路径，比如 epoch_5.pt
)
@torch.no_grad() #这个函数不需要梯度，用装饰器关闭梯度计算，节省显存和计算。
def extract_embedding(model, img_path, transform, device):#定义提取单张图片embedding 的函数
img = Image.open(img_path).convert("RGB") # 打开图片并转换成 RGB 格式
x = transform(img).unsqueeze(0).to(device) #预处理图片，增加batch维度，并移动到CPU/GPU
emb = model(x) 通过embedding 网络提取特征
emb = F.normalize(emb, dim=1)对输出特征做L2归一化，便于计算余弦相似度
return emb[0]#返回batch中第1个，也就是这张图片的embedding
def match(args):#定义两张图片的人脸匹配函数
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") #如果有GPU就是用cuda, 否则用cpu
ckpt = torch.laod(args.ckpt, map_location=device) 加载训练好的checkpoint并映射到当前设备
model = EmbeddingNet(emb_dim=ckpt["emb_dim"]).to(device) #按checkpoint中的维度重建embedding网络
model.load_state_dict(ckpt["backbone"]) #加载训练好的backbone参数
model.eval() 设置为评估模式，关闭dropout 固定BatchNorm行为
transform = build_eval_transform(ckpt["img_size"])#创建推理阶段图片预处理流程
emb1 = extract_embedding(mode, args.img1, transform, device) #提取第一张图片的人脸embedding
emb2 = extract_embedding(model, args.img2, transform, device)
cosine_sim = torch.dot(emb1, emb2).item() 两个归一化embedding点积，等价于余弦相似度
def main(): 程序入口函数
parser = argparse.ArgumentParser() # 创建命令行参数解析器
subparsers = parser.add_subparsers(dest="command", required=True) # 创建子命令，比如 train 和 match
train_parser = subparsers.add_parser("train") # 添加 train 子命令
train_parser.add_argument("--data", type=str, required=True) # 训练数据路径，必须提供
train_parser.add_argument("--out", type=str, default="runs/face_arcface") # 模型输出目录
train_parser.add_argument("--epochs", type=int, default=30) # 训练 epoch 数
train_parser.add_argument("--batch-size", type=int, default=128) # batch size
train_parser.add_argument("--workers", type=int, default=4) # DataLoader 子进程数量
train_parser.add_argument("--img-size", type=int, default=112) # 输入图片大小
train_parser.add_argument("--emb-dim", type=int, default=512) # embedding 特征维度
train_parser.add_argument("--lr", type=float, default=1e-3) # 学习率
train_parser.add_argument("--weight-decay", type=float, default=5e-4) # AdamW 权重衰减
train_parser.add_argument("--margin", type=float, default=0.5) # ArcFace margin
train_parser.add_argument("--scale", type=float, default=64.0) # ArcFace scale
train_parser.add_argument("--save-every", type=int, default=5) # 每隔多少个 epoch 保存一次模型
train_parser.add_argument("--seed", type=int, default=42) # 随机种子
train_parser.add_argument("--amp", action="store_true") # 是否启用混合精度训练
match_parser = subparsers.add_parser("match") # 添加 match 子命令
match_parser.add_argument("--ckpt", type=str, required=True) # checkpoint 路径，必须提供
match_parser.add_argument("--img1", type=str, required=True) # 第一张待匹配图片路径
match_parser.add_argument("--img2", type=str, required=True) # 第二张待匹配图片路径
match_parser.add_argument("--threshold", type=float, default=None) # 匹配阈值，不传则只输出相似度
args = parser.parse_args() # 解析命令行参数
if args.command == "train": # 如果用户运行的是 train 子命令
train(args) # 调用训练函数
elif args.command == "match": # 如果用户运行的是 match 子命令
match(args) # 调用匹配函数
if __name__ == "__main__": # 只有直接运行当前文件时才执行 main
main() # 调用程序入口net.conv1 = nn.Conv2d(
3,
64,
kernel_size=3,
stride=1,
padding=1,
bias=False
)
把7x7卷积改成3x3 stride=1, 分辨率不要降低
2.2.3 maxpool移除
net.maxpool = nn.Identity() 直接透传
2.3 EmbeddingNet实际网络结构
经过网络后，大致结构
输入 RGB人脸图像， 输出尺寸 [B,3,112,112]
conv1 3x3卷积，64通道,stride = 1， 输出 [B, 64, 112x112]
maxpool Identity() 输出 [B, 64, 112x112]
layer1 ResNet BasicBlockx2, [B, 64, 112, 112]
layer2 ResNet BaisicBlock x2 下采样， [B, 128, 56, 56]
layer3 ResNet BasicBlock x2 下采样， [B, 256, 28, 28]
layer4 ResNet BasicBlockx2 下采样 [B,512, 14, 14]
avgpool 全局平均池化， [B, 512, 1, 1]
flatten 展平， [B, 512]
fc 全连接 [B, 512]
BatchNorm1d embedding归一化稳定训练， [B, 512]
2.4 ResNet18厘米的残差结构
普通CNN x->卷积 -> F(x)
ResNet
x->卷积层->F(x) ->F(x) + x
也就是每个block输出，
输出 = 卷积分支结果 + 原始输入
2.5 ArcFace分类头结构
class ArcMarginProduct(nn.Module):
不是CNN，是一个训练用的分类头
1 对embedding做L2 normalize
2 对每个类别的分类权重做L2 normalize
3 计算embedding和每个类别权重的Cosine相似度
4 乘以scale
6 输入CrossEntropyLoss
普通分类头是
embedding->Linear->logits->CrossEntropy
ArcFace是
embedding-<normalize->cosine->加角度margin->scale->CrossEntropy
目的不是单纯分类，让特征空间更适合人脸匹配
2.6 训练时输入数据格式
datasets.ImageFloder(args.data)
训练集目录必须是这种格式
data/train/person_001/xxx.jpg
data/train/person_002/xxx.jpg
data/train/person_003/xxx.jpg
2.7 训练图像预处理流程
build_train_transform(args.img_size)
默认 img_size = 112
人脸已经裁剪好，眼睛基本对齐，脸在图片中心，背景尽量少
2.7.2 随机水平翻转
transforms.RandomHorizontalFlip(p=0.5)
增强模型鲁棒性
2.7.3 随机放射变换
transforms.RandomAffine(
degrees = 5,
translate = (0.03, 0.03),
scale=(0.95, 1.05)
)
2.7.4 颜色扰动
transforms.ColorJitter(
brightness = 0.15
contrast=0.15
saturation=0.10
hue = 0.03
) 对图片做一些随机调整
亮度 brightness
对比度 contrast
饱和度 saturation
色相 hue
2.7.5 转成tensor 张量
transforms.ToTensor()
PIL图片原本像素范围一般是0~255
经过ToTensor 后会变成
0.0 ~ 1.0
2.7.6 Normalize归一化
transforms.Normalize(
0.5, 0.5, 0.5\], \[0.5, 0.5, 0.5
)
这一步对RGB三个通道分别做
新像素值 = （原始像素 - 0.5）/0.5
2.8 训练阶段完整数据流
原始图片
->resize 112x112
->随机翻转
->随机旋转，【平移，缩放
->随机颜色扰动
->ToTensor 变成[3,112,112]
->Normalize 归一化 [-1, 1]
组成batch [B,3,112,112] 几张图一组
->EmbeddingNet
->512 维embedding
->ArcFace head
->num_class 维logits
-〉 crossEntropyLoss
-》反向传播更新模型
2.9 推理/匹配阶段预处理
推理阶段使用
build_eval_transform()
比训练阶段简单
transforms.Resize(img_size, img_size)
transforms.ToTensor()
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
推理时需要稳定结果，不能随机翻转，随机颜色变化
图A
->Resize + ToTensor + Normalize
->EmbeddingNet
->embedding A
L2 normalize
图片B
->Resize + ToTensor + normalize
->EmbeddingNet
->embedding B
->L2 normalize
cosine_similarity = embedding A - embedding B
如果相似度大于阈值，就认为是同一个人。
2.10 这个模型训练的本质
训练时 输入 单张人脸图像，这个人是谁，身份分类
模型通过身份分类学到一个好的特征空间
训练完成后
同一个人的embedding会靠近，不同人的embedding会远离
2.11 训练前你最好额外做的预处理
1 人脸检查，人脸裁剪，人脸关键点检测，根据双眼位置做人脸对齐，保存成112x112的人脸图
理性训练图片应该类似
整张图主要是人脸，双眼大致水平，人脸位于中心，背景尽量少，尺度比较统一
2.12 总结
改造ResNet18 CNN作为人脸提取器
使用ArcFace分类进行身份分类训练
训练时输入单张人脸图片和身份标签
推理时只使用CNN提取512维embedding
最后通过余弦相似度判断两张人脸是否属于同一个人