十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

SadTalker 内置 ArcFace 人脸识别训练基准:Partial FC 与 Model Parallel 的训练速度与显存对比实测

SadTalker 内置 ArcFace 人脸识别训练基准:Partial FC 与 Model Parallel 的训练速度与显存对比实测 SadTalker 内置 ArcFace 人脸识别训练基准Partial FC 与 Model Parallel 的训练速度与显存对比实测【免费下载链接】SadTalker[CVPR 2023] SadTalkerLearning 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_classes300 万合成数据synthetic数据来源合成数据config.rec synthetic混合精度训练开启config.fp16 True主干网络backboneResNet50配置中config.network r50单卡 batch size1024文档描述配置文件中为 128 × 8 卡 1024 全局两个测试用例的配置分别位于 3millions.pyModel Parallelsample_rate 1.0与 3millions_pfc.pyPartial FCsample_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_rate1.0 python -m torch.distributed.launch --nproc_per_node8 --nnodes1 --node_rank0 --master_addr127.0.0.1 --master_port1234 train.py configs/3millions # Partial FC 0.110% 类中心采样sample_rate0.1 python -m torch.distributed.launch --nproc_per_node8 --nnodes1 --node_rank0 --master_addr127.0.0.1 --master_port1234 train.py configs/3millions_pfc参数含义--nproc_per_node8表示每节点 8 个进程对应 8 张 GPU--nnodes1 --node_rank0表示单机训练--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_norm5的梯度裁剪损失函数由config.loss arcface指定losses.py 中的ArcFace模块在余弦相似度上施加角度裕量m0.5并乘以尺度s64.0。三、Partial FC 的底层原理源码级解读要理解为什么 Partial FC 能同时降低显存并提升吞吐需要看 partial_fc.py 的核心实现。初始化时每个 rankGPU只持有本地分片的类中心partial_fc.pyself.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_reduceReduceOp.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 MBPartial FC 0.1单卡占用约10.3~10.5 GB8 张卡均在 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/secondTesla V100 32GB × 8越大越好数据集身份数量Data ParallelModel ParallelPartial FC 0.1125000468148245004250000404745214976500000308740134900100000020903449480314000001672304347382000000-259346264000000-174842085500000-138939758000000--356516000000--267929000000--1855-表示该方案在此身份规模下因显存不足或训练过慢而未测。GPU 显存开销对比GB/GPUTesla V100 32GB × 8越小越好数据集身份数量Data ParallelModel ParallelPartial FC 0.11250007358530648682500009940582650045000001422071145202100000023708996656201400000322521117860562000000-1397864724000000-2323882845500000-3218898548000000--1231016000000--1995029000000--32324-表示超出单卡 32GB 显存限制或未测。4.3 数据规律解读从上表可以归纳出三个工程判断身份规模较小时差距不大12.5 万身份时三种方案吞吐都在 4700~5000 samples/sec 之间显存差异也不大此时用最简单的 Data Parallel 即可。身份规模超过百万后分化显著100 万身份时 Data Parallel 吞吐已降至 2090 samples/sec且单卡显存冲到 23.7GB而 Model Parallel3449与 Partial FC 0.14803仍保持高吞吐到 140 万身份时 Data Parallel 显存已达 32.2GB触及单卡上限。Partial FC 0.1 是超大规模的唯一可行路径Model Parallel 到 550 万身份时显存已逼近 32GB32188 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] SadTalkerLearning Realistic 3D Motion Coefficients for Stylized Audio-Driven Single Image Talking Face Animation项目地址: https://gitcode.com/GitHub_Trending/sa/SadTalker创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表