SadTalker 内置 ArcFace 人脸识别训练基准:Partial FC 与 Model Parallel 的训练速度与显存对比实测
【免费下载链接】SadTalker[CVPR 2023] SadTalker:Learning Realistic 3D Motion Coefficients for Stylized Audio-Driven Single Image Talking Face Animation项目地址: https://gitcode.com/GitHub_Trending/sa/SadTalker
导读
本文以仓库内置的 speed_benchmark.md 基准测试文档为核心,深入讲解 ArcFace 人脸识别训练中两种大规模分类并行策略——**Model Parallel(模型并行)**与Partial FC 0.1(部分全连接采样)——的实测性能对比。文档配套的测试命令、配置与数据均可在 arcface_torch 目录下复现。读完本文,你将掌握:如何复现该基准测试、两份基准配置(3millions.py 与 3millions_pfc.py)的差异本质,以及在不同类别规模(12.5 万至 2900 万身份)下训练吞吐与单卡显存的变化规律,为大规模人脸识别训练的资源选型提供量化依据。
说明:本文涉及的训练代码位于 train.py,核心并行实现位于 partial_fc.py。所有速度与显存数据均来自关联文档,属于仓库作者在特定硬件环境(8×Tesla V100 32GB)下的实测记录。
一、基准测试背景:为什么需要对比这两种并行方式
人脸识别模型(ArcFace)训练的本质,是在特征嵌入空间中让同一身份(identity/class)的样本聚拢、不同身份的样本分离。最终分类层的权重矩阵维度为embedding_size × num_classes,其中num_classes是数据集中身份数量。当身份数量达到数百万甚至数千万级别时,分类层权重本身就会占据巨大的显存与计算量,成为训练瓶颈。
仓库在 train.py 中提供了两种应对思路:
- Model Parallel(模型并行):将完整的分类层权重按身份切分到多张 GPU 上,每张卡只负责自己那部分类中心,配合
all_gather/all_reduce跨卡同步特征与梯度; - Partial FC(部分全连接):在模型并行的基础上,每步训练只采样一小部分类中心参与 Softmax 计算(通过
sample_rate控制),大幅削减分类层的显存与计算开销。其论文出处为 Partial FC: Training 10 Million Identities on a Single Machine。
基准文档通过完全相同的硬件与数据集规模,对比这两种方案(以及作为对照的纯数据并行 Data Parallel)在吞吐和显存上的表现,从而回答一个工程问题:身份数量多大时,必须从数据并行切换为模型并行 / Partial FC?
二、测试环境与复现命令
2.1 测试条件
关联文档明确给出了测试条件:
| 项目 | 设定 |
|---|---|
| 身份数量(num_classes) | 300 万(合成数据,synthetic) |
| 数据来源 | 合成数据(config.rec = "synthetic") |
| 混合精度训练 | 开启(config.fp16 = True) |
| 主干网络(backbone) | ResNet50(配置中config.network = "r50") |
| 单卡 batch size | 1024(文档描述),配置文件中为 128 × 8 卡 = 1024 全局 |
两个测试用例的配置分别位于 3millions.py(Model Parallel,sample_rate = 1.0)与 3millions_pfc.py(Partial FC,sample_rate = 0.1),两者除sample_rate外其余参数完全一致:
config.loss = "arcface" # 损失函数:ArcFace config.network = "r50" # 主干网络:ResNet50 config.embedding_size = 512 # 特征嵌入维度 config.sample_rate = 1.0 / 0.1 # 类中心采样率(唯一差异点) config.fp16 = True # 开启混合精度训练 config.batch_size = 128 # 单卡 batch size config.lr = 0.1 # 基础学习率 config.rec = "synthetic" # 使用合成数据,避免 I/O 干扰 config.num_classes = 300 * 10000 # 300 万身份 config.num_epoch = 30从源码可以确认,sample_rate = 1.0时 PartialFC 退化为完整分类层(self.update = lambda: 0,直接持有全量权重);只有sample_rate < 1.0时才会走采样分支(sample方法),这正解释了为什么两个测试命令跑出的结果差异完全可归因于采样率。
2.2 复现命令
文档给出两条测试命令(8 卡、单机、NCCL 分布式):
# Model Parallel(全量类中心,sample_rate=1.0) python -m torch.distributed.launch --nproc_per_node=8 --nnodes=1 --node_rank=0 --master_addr="127.0.0.1" --master_port=1234 train.py configs/3millions # Partial FC 0.1(10% 类中心采样,sample_rate=0.1) python -m torch.distributed.launch --nproc_per_node=8 --nnodes=1 --node_rank=0 --master_addr="127.0.0.1" --master_port=1234 train.py configs/3millions_pfc参数含义:--nproc_per_node=8表示每节点 8 个进程(对应 8 张 GPU);--nnodes=1 --node_rank=0表示单机训练;--master_addr/--master_port指定分布式主节点地址与端口;train.py后的第一个位置参数即配置文件路径(不含.py后缀)。
在 train.py 中可以看到,当cfg.rec == "synthetic"时,训练数据由 SyntheticDataset 提供——它固定生成一张 112×112 的随机噪声图像与标签1,数据集长度为 1,000,000。合成数据的意义在于排除真实数据集的磁盘 I/O 与解码开销,使基准测试结果纯粹反映计算与通信性能。
2.3 训练入口的关键调用链
基准测试的运行路径如下(对应 train.py 的训练循环):
backbone = get_model(cfg.network, ...)构建 ResNet50 主干,network = "r50"在 backbones/init.py 中映射到iresnet50;- 特征归一化后调用
module_partial_fc.forward_backward(label, features, opt_pfc)完成分类层的前向与反向(见 partial_fc.py); - 在
fp16模式下,特征梯度经MaxClipGradScaler缩放后再反传到主干,并对主干参数做max_norm=5的梯度裁剪; - 损失函数由
config.loss = "arcface"指定,losses.py 中的ArcFace模块在余弦相似度上施加角度裕量m=0.5并乘以尺度s=64.0。
三、Partial FC 的底层原理(源码级解读)
要理解为什么 Partial FC 能同时降低显存并提升吞吐,需要看 partial_fc.py 的核心实现。初始化时,每个 rank(GPU)只持有本地分片的类中心(partial_fc.py):
self.num_local = num_classes // world_size + int(rank < num_classes % world_size) self.class_start = num_classes // world_size * rank + min(rank, num_classes % world_size) self.num_sample = int(self.sample_rate * self.num_local)即 300 万身份在 8 卡上每卡本地持有 37.5 万个类中心,Partial FC 0.1 再从其中采样约 3.75 万个参与本轮 Softmax。
训练时,prepare 方法执行跨卡all_gather汇总全局标签,sample方法(partial_fc.py)做两件事:
- 保留正样本:将属于本卡负责区间内的标签映射为局部索引;
- 随机补足负样本:用
torch.topk在均匀随机数中选取num_sample个索引,保证正样本类中心必定被选中,其余位置由随机负样本填充。
随后 forward_backward 通过all_gather聚合全局特征、计算采样后的 logits,在no_grad块中用all_reduce(ReduceOp.MAX与ReduceOp.SUM)完成跨卡数值稳定的 Softmax 归一化,最后用reduce_scatter把特征梯度分发回各卡。这一整套流程实现了「全局 Softmax 精确计算 + 每步仅计算 10% 类中心」的效果——显存只存采样子集,算力只花在采样子集上。
值得一提的是,当sample_rate = 1.0时(Model Parallel 场景),partial_fc.py 直接走self.update = lambda: 0的旁路,sub_weight即为全量本地权重,不执行任何采样。
四、基准测试实测数据
以下数据全部来自 speed_benchmark.md 的原文记录。
4.1 300 万身份实测:显存与速度
GPU 显存占用(gpustat -i采样,单位 MB,单卡 32510 MB 总量):
- Model Parallel:单卡占用约28.8~30.3 GB(其中 0 号卡 30338 MB,其余卡约 28872~28932 MB);
- Partial FC 0.1:单卡占用约10.3~10.5 GB(8 张卡均在 10328~10488 MB 区间)。
训练吞吐(training.log中Speed字段,单位 samples/sec):
# (Model Parallel) training.log Training: Speed 2271.33 samples/sec Loss 1.1624 LearningRate 0.2000 Epoch: 0 Global Step: 100 Training: Speed 2269.94 samples/sec Loss 0.0000 LearningRate 0.2000 Epoch: 0 Global Step: 150 Training: Speed 2272.67 samples/sec Loss 0.0000 LearningRate 0.2000 Epoch: 0 Global Step: 200 # (Partial FC 0.1) training.log Training: Speed 5299.56 samples/sec Loss 1.0965 LearningRate 0.2000 Epoch: 0 Global Step: 100 Training: Speed 5296.37 samples/sec Loss 0.0000 LearningRate 0.2000 Epoch: 0 Global Step: 150 Training: Speed 5304.37 samples/sec Loss 0.0000 LearningRate 0.2000 Epoch: 0 Global Step: 200文档的结论非常明确:在此测试用例中,Partial FC 0.1 的显存占用仅为 Model Parallel 的约 1/3,训练速度则快约 2.5 倍(吞吐由约 2270 samples/sec 提升至约 5300 samples/sec)。
4.2 不同身份规模下的横向对比
训练速度对比(samples/second,Tesla V100 32GB × 8,越大越好):
| 数据集身份数量 | Data Parallel | Model Parallel | Partial FC 0.1 |
|---|---|---|---|
| 125000 | 4681 | 4824 | 5004 |
| 250000 | 4047 | 4521 | 4976 |
| 500000 | 3087 | 4013 | 4900 |
| 1000000 | 2090 | 3449 | 4803 |
| 1400000 | 1672 | 3043 | 4738 |
| 2000000 | - | 2593 | 4626 |
| 4000000 | - | 1748 | 4208 |
| 5500000 | - | 1389 | 3975 |
| 8000000 | - | - | 3565 |
| 16000000 | - | - | 2679 |
| 29000000 | - | - | 1855 |
(-表示该方案在此身份规模下因显存不足或训练过慢而未测。)
GPU 显存开销对比(GB/GPU,Tesla V100 32GB × 8,越小越好):
| 数据集身份数量 | Data Parallel | Model Parallel | Partial FC 0.1 |
|---|---|---|---|
| 125000 | 7358 | 5306 | 4868 |
| 250000 | 9940 | 5826 | 5004 |
| 500000 | 14220 | 7114 | 5202 |
| 1000000 | 23708 | 9966 | 5620 |
| 1400000 | 32252 | 11178 | 6056 |
| 2000000 | - | 13978 | 6472 |
| 4000000 | - | 23238 | 8284 |
| 5500000 | - | 32188 | 9854 |
| 8000000 | - | - | 12310 |
| 16000000 | - | - | 19950 |
| 29000000 | - | - | 32324 |
(-表示超出单卡 32GB 显存限制或未测。)
4.3 数据规律解读
从上表可以归纳出三个工程判断:
- 身份规模较小时差距不大:12.5 万身份时三种方案吞吐都在 4700~5000 samples/sec 之间,显存差异也不大;此时用最简单的 Data Parallel 即可。
- 身份规模超过百万后分化显著:100 万身份时 Data Parallel 吞吐已降至 2090 samples/sec,且单卡显存冲到 23.7GB,而 Model Parallel(3449)与 Partial FC 0.1(4803)仍保持高吞吐;到 140 万身份时 Data Parallel 显存已达 32.2GB,触及单卡上限。
- Partial FC 0.1 是超大规模的唯一可行路径:Model Parallel 到 550 万身份时显存已逼近 32GB(32188 MB)且吞吐降至 1389;而 Partial FC 0.1 可以一路扩展到 2900 万身份(显存 32324 MB、吞吐 1855 samples/sec),单机 8 卡即可训练千万级身份的人脸识别模型。
五、在 SadTalker 项目中的上下文与可复现性
该基准测试代码随 SadTalker 仓库以子模块形式内置,位于 src/face3d/models/arcface_torch。需要说明的是:ArcFace 训练基准与 SadTalker 的推理主流程无直接耦合——SadTalker 利用的是该子模块中已经训练好的 ArcFace 人脸识别特征提取能力(相关依赖见 requirements3d.txt),用于人脸关键点检测与三维人脸重建中的身份特征提取;而speed_benchmark.md记录的训练基准服务于希望从零训练或继续训练大规模人脸识别模型的进阶用户。
若要复现该基准,需要:
- 8×Tesla V100 32GB 或等价显存的 GPU 集群(单机即可);
- PyTorch 与 NCCL 分布式环境,混合精度(AMP)支持;
- 运行本文第 2.2 节的两条命令,观察
training.log中的Speed字段与gpustat -i的显存占用。
对显存敏感、又没有多卡集群的读者,Partial FC 的采样思路(partial_fc.py)本身也可作为单卡训练千万级分类任务的参考实现。
六、总结
围绕 speed_benchmark.md 的实测数据可以得出:身份规模在百万级以下时三种并行方案差距不大,百万级以上 Model Parallel 优于 Data Parallel,而 Partial FC 0.1 在显存与吞吐两个维度上全面领先,并能支撑千万级身份的分布式训练。配套的两份基准配置 3millions.py 与 3millions_pfc.py 仅相差一个sample_rate参数(1.0 vs 0.1),其背后的采样、全归约与梯度分发逻辑均可从 partial_fc.py 与 train.py 源码中一一印证——这也让本文的数据具有可复现、可追溯的技术依据。
【免费下载链接】SadTalker[CVPR 2023] SadTalker:Learning Realistic 3D Motion Coefficients for Stylized Audio-Driven Single Image Talking Face Animation项目地址: https://gitcode.com/GitHub_Trending/sa/SadTalker
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考