一 代码分析
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
最后通过余弦相似度判断两张人脸是否属于同一个人