InsightFace Partial FC 人脸识别大规模分布式训练框架与 Glint360K 数据集完全指南
【免费下载链接】insightfaceState-of-the-art 2D and 3D Face Analysis Project项目地址: https://gitcode.com/GitHub_Trending/in/insightface
Partial FC(Partial Face Classification)是 InsightFace 仓库中面向大规模人脸识别分类任务(例如 1000 万、1 亿身份级别)的分布式深度学习训练框架。本文以其官方文档为核心,结合仓库内 MXNet 实现源码(train_memory.py、memory_module.py、memory_bank.py)与配套配置,系统讲解 Partial FC 的原理、环境搭建、单机/多机训练、Glint360K 数据集的下载解压与评测结果,帮助读者掌握如何在单机 8 卡乃至 64 卡环境下训练千万级身份的人脸识别模型。
一、Partial FC 是什么
Partial FC 的定位是一个面向人脸识别的大规模分布式训练框架,核心目标是解决大规模分类任务(例如 1000 万或 1 亿个身份类别)训练时的两大痛点:
- 显存瓶颈:常规 Softmax 分类层需要维护一张
类别数 × 特征维度的全量权重矩阵,身份数量达到千万级后,仅分类层权重就远超单卡显存; - 通信开销:传统 Model Parallel(模型并行)方案需要频繁同步全部类别中心,吞吐受限。
Partial FC 的解决方案是:分类层(类别中心)按 GPU 均分存储 + 负类中心随机采样近似 Softmax。官方文档明确指出,它比模型并行方案快得多,且无精度损失(no performance drop)。从源码结构看,Partial FC 的训练主体由以下模块协作完成:
| 模块文件 | 职责 |
|---|---|
| train_memory.py | 训练入口,基于 horovod 初始化、组装 MemoryBank 与 MarginLoss |
| memory_module.py | 核心训练模块SampleDistributeModule,实现特征 AllGather、Softmax 分母 AllReduce、梯度回传 |
| memory_bank.py | MemoryBank,负责本 rank 局部类别中心及其动量(momentum)的存取与采样 |
| memory_softmax.py | MarginLoss,ArcFace / CosFace 间隔损失的局部实现 |
| memory_samplers.py | 负类中心采样器 |
| default.py | 数据集 / 网络 / 损失等全部超参配置 |
二、算法原理:分类层模型并行与 Softmax 近似
官方文档将 Partial FC 的方法归纳为两大部分,仓库源码提供了对应的实现证据。
1. 分类层的模型并行
类别中心被均匀分布到不同的 GPU 上,仅需三次通信即可完成一次无损(loss-free)的 Softmax 计算。以 memory_module.py 中的实现为例:
① 特征的同步(AllGather)要保证每块 GPU 上都持有全部 GPU 的特征。forward()中调用allgather(),其内部通过 in-place AllReduce 实现,将本 rank 的 batch 写入对应位置后做全量求和:
total_tensor[self.rank * self.batch_size:self.rank * self.batch_size + self.batch_size] = tensor hvd.allreduce_(total_tensor, average=False) # all-reduce in-place② Softmax 分母的同步(AllReduce)先在本地计算exp(logits)的局部和,再通过通信得到全局和。backward()与backward_sample()中,global_sum_fc7 = hvd.allreduce(sum_fc7, average=False)即对应文档中Allreduce(sum(exp(logits_i)))这一步;同时在求exp之前先同步全局最大值做数值稳定处理(global_max_fc7)。
③ 特征梯度的同步(AllReduce)logits 的梯度可以独立计算,特征的梯度自然也能独立得到,最后将全部 GPU 上的梯度收集并回传给 backbone:
total_feature_grad = hvd.allreduce(total_feature_grad, average=False) fc1_grad = total_feature_grad[self.batch_size * self.rank:self.batch_size * self.rank + self.batch_size] self.backbone_module.backward(out_grads=[fc1_grad / self.size])2. Softmax 的近似
只需类别中心的一个子集即可近似完整 Softmax 的计算,但正类中心必须包含在被采样集合中。官方文档给出了伪代码:
centers_p = func_positive(label) # select the positive class centers by the label of the sample centers_n = func_negative(centers_p) # negative class centers are randomly sampled after excluding positive classes centers_final = concat(centers_n, centers_p) # class centers that participate in softmax calculations在 memory_bank.py 中,采样由WeightIndexSampler完成:sample()先对全局 label 做 unique 与 sort,再由weight_index_sampler(global_label)得到局部索引,返回的索引经get()取出对应的类别中心与动量,参与当次前向。是否采样由配置项sample_ratio控制:backward_all()中if not bool(config.sample_ratio - 1):走全量backward(),否则走backward_sample()。
MarginLoss 与主流损失兼容:memory_softmax.py 的MarginLoss默认实现 ArcFace,并通过loss_m1/loss_m2/loss_m3三元组兼容 CosFace(m1=1.0, m2=0.0时退化为 CosFace,m3为余弦间隔),对特征与权重分别做 L2 归一化后点积得到fc7,最后乘缩放因子loss_s(默认 64)。
三、配套数据集 Glint360K
文档将 Partial FC 与Glint360K数据集捆绑发布:该数据集经清洗、合并后共包含17,091,657 张图像、360,232 个身份,是目前最大且最干净的人脸识别数据集之一。使用 Partial FC 策略在 Glint360K 上训练的基线模型即可达到当时的最优水平。
1. IFRT 评测(按人种细分)
下表为 IFRT 大型测试集上的评估结果,其中r表示负类中心采样率(sampling rate of negative class centers),这也是 Partial FC 最重要的超参数:
| Backbone | Dataset | African | Caucasian | Indian | Asian | ALL |
|---|---|---|---|---|---|---|
| R50 | MS1M-V3 | 76.24 | 86.21 | 84.44 | 37.43 | 71.02 |
| R124 | MS1M-V3 | 81.08 | 89.06 | 87.53 | 38.40 | 74.76 |
| R100 | Glint360k(r=1.0) | 89.50 | 94.23 | 93.54 | 65.07 | 88.67 |
| R100 | Glint360k(r=0.1) | 90.45 | 94.60 | 93.96 | 63.91 | 88.23 |
可以看到,Glint360K 在 IFRT 全部人种子集上全面超越 MS1M-V3,且采样率r=0.1时部分指标(African/Caucasian/Indian)甚至优于r=1.0,这印证了负类采样在大规模分类任务中无损甚至更优的特性。
2. IJB-C 与 MegaFace 评测
评测配置为 ResNet100 backbone + CosFace(m=0.4)损失函数;IJB-C 报告 TAR@FAR=1e-4,MegaFace 报告 TAR@FAR=1e-6:
| Test Dataset | IJB-C | Megaface_Id | Megaface_Ver |
|---|---|---|---|
| MS1MV2 | 96.4 | 98.3 | 98.6 |
| Glint360k | 97.3 | 99.1 | 99.1 |
3. 许可证
Glint360K 数据集(及其上训练的模型)仅限非商业研究用途(non-commercial research purposes only)。
4. 下载与解压
提供百度网盘(提取码o3az)与磁力链接两种渠道,磁力 URI 为:
magnet:?xt=urn:btih:E5F46EE502B9E76DA8CC3A0E4F7C17E4000C7B1E&dn=glint360k分卷压缩包使用如下命令解压(注意cat管道末尾的-不能漏掉):
cat glint360k_* | tar -xzvf - # Don't forget the last '-'!各分卷的 MD5 校验值(用于核对下载完整性):
# cf7433cbb915ac422230ba33176f4625 glint360k_00 # 589a5ea3ab59f283d2b5dd3242bc027a glint360k_01 # 8d54fdd5b1e4cd55e1b9a714d76d1075 glint360k_02 # cd7f008579dbed9c5af4d1275915d95e glint360k_03 # 64666b324911b47334cc824f5f836d4c glint360k_04 # a318e4d32493dd5be6b94dd48f9943ac glint360k_05 # c3ae1dcbecea360d2ec2a43a7b6f1d94 glint360k_06解压得到的.rec/.idx记录文件 MD5:
# 5d9cd9f262ec87a5ca2eac5e703f7cdf train.idx # 8483be5af6f9906e19f85dee49132f8e train.rec如需还原为图片,使用仓库中的 unpack_glint360k.py 进行解包。
5. 预训练模型
预训练模型提供百度网盘(提取码befi)与 Google Drive 下载。各模型关键指标如下(IJBC@e4 即 TAR@FAR=1e-4,IFRT@e6 即 TAR@FAR=1e-6):
| Framework | backbone | negative class centers sample_rate | IJBC@e4 | IFRT@e6 |
|---|---|---|---|---|
| mxnet | R100 | 1.0 | 97.3 | - |
| mxnet | R100 | 0.1 | 97.3 | - |
| pytorch | R50 | 1.0 | 97.0 | - |
| pytorch | R100 | 1.0 | 97.4 | - |
6. 数据集 FAQ
- 对齐设置:Glint360K 采用与 MS1MV2 相同的对齐方式;
- 为何更新 Glint360K(旧版是否有 Bug):旧版本在使用 Softmax 训练时无问题,但在 triplet 训练中存在 Bug,最新版已修复;
- 是否有 Google Drive 或 Dropbox 版本:已发布 torrent(磁力链接)。
四、Docker 环境:开箱即用的训练镜像
文档为 Partial FC 提供了官方 Docker 镜像,避免了手工编译 horovod/mxnet 的繁琐过程。
1. 拉取镜像
docker pull insightface/partial_fc:v1(离线 docker.tar 镜像文档标注 "coming soon",以 Docker Hub 拉取为准。)
2. 启动容器
sudo docker run -it -v /train_tmp:/train_tmp --net=host --privileged --gpus 8 --shm-size=1g insightface/partial_fc:v1 /bin/bash其中/train_tmp是存放训练集的目录(如果机器内存足够大,可先将其挂载为tmpfs内存盘)。注意镜像内 CUDA 版本为10.1,因此物理机显卡驱动版本必须大于 418;宿主机器上无需安装 CUDA Toolkit,但必须安装 NVIDIA 驱动。
五、源码级安装与训练(MXNet 版)
除 Docker 外,mxnet/README.md 提供了从零搭建环境的完整流程。
1. 环境要求
| 组件 | 版本 |
|---|---|
| python | ==3.6 |
| cuda | ==10.1 |
| cudnn | ==765 |
| mxnet-cu101 | ==1.6.0.post0 |
| nccl | 推荐安装(有则更快) |
| openmpi | ==4.0.0(需源码编译,见 install-mpi.sh) |
| horovod | ==0.19.2(见 install-horovod.sh) |
Python 依赖安装:
pip install easydict mxboard opencv-python tqdm版本兼容性关键提示:部分 mxnet 版本无法安装 horovod,官方强烈建议mxnet==1.6.0 + cuda==10.1。社区已知mxnet 1.5.1 无法安装 horovod,建议尝试 mxnet 1.5 或 1.6 版本。
2. 启动训练(horovodrun 与 mpirun)
horovod 底层仍调用 MPI:有多少块 GPU 就启动多少个进程,进程数通过-np指定。单机 8 卡:
horovodrun -np 8 -H localhost:8 bash config.sh两台机器共 16 卡:
horovodrun -np 16 -H ip1:8,ip2:8 bash config.sh或直接使用 mpirun:
bash run.sh多机训练前,horovodrun所在主机必须能免密 SSH到所有其他主机(含自身),可执行ssh-copy-id user@ip配置。
3. 训练入口与配置解析
config.sh 是实际的训练命令,它设置 NCCL 后端环境变量后调用 train_memory.py:
export CUDA_VISIBLE_DEVICES='0,1,2,3,4,5,6,7' export HOROVOD_GPU_ALLREDUCE=NCCL export HOROVOD_GPU_ALLGATHER=NCCL export HOROVOD_GPU_BROADCAST=NCLL export MXNET_CPU_WORKER_NTHREADS=3 PYTHON_EXEC=/usr/bin/python ${PYTHON_EXEC} train_memory.py \ --dataset glint360k_8GPU \ --loss cosface \ --network r100 \ --models-root /data/anxiang/opensource/glint360k_8GPU_r100FC_1.0_fp32_cosface三个必选参数分别对应 default.py 中的配置族:
--dataset:预置数据集配置。常用值包括glint360k_8GPU(num_classes=360232、batch_size=64、max_update=600000、lr_steps=200000,400000,500000,550000)、glint360k_16GPU、emore(85742 类)、webface(10575 类),以及用于调试/压力测试的100w/1000w/2000w/3000w/10000w(这些调试配置会打开config.debug = 1);--loss:cosface(m3=0.4)或arcface(m2=0.5);--network:r100、r122等(设置net_name=resnet与num_layers)。
default.py 中的核心全局配置:
config.embedding_size = 512 # 特征维度 config.image_size = 112 # 输入图像尺寸 config.batch_size = 64 # 每卡 batch size config.backbone_lr = 0.1 # backbone 初始学习率 config.memory_bank_lr = config.backbone_lr # 类别中心(memory bank)学习率 config.sample_ratio = 1.0 # 负类中心采样率 r,核心超参数 config.fp16 = False # 是否混合精度 config.debug = 0 # 置 1 时使用 DummyIter 屏蔽 IO4. 训练流程中的关键实现细节
在 train_memory.py 中:
- 类别中心(即 Softmax 线性变换矩阵)按 rank等分存储:
num_local = (num_classes + size - 1) // size,每卡实际参与分类的采样数num_sample = int(num_local * sample_ratio); MemoryBank以 GPU 存储类别中心weight与动量weight_mom(均以N(0, 0.01)初始化),并提供sample / get / set / save四个核心方法,训练中通过CallBackCenterSave周期性保存各 rank 的*_centers.param与*_centers_mom.param;- backbone 与 memory bank 使用两个独立优化器:backbone 用
DistributedOptimizer(SGD(...)),momentum=0.9、wd=5e-4,且rescale_grad=1.0 / (config.batch_size * size) * size;memory bank 用MemoryBankSGDOptimizer,rescale_grad=1.0 / config.batch_size / size; - 文件开头设置了一批针对 MXNet + horovod 的性能环境变量,例如
HOROVOD_FUSION_THRESHOLD=67108864、HOROVOD_NUM_NCCL_STREAMS=2、MXNET_EXEC_BULK_EXEC_MAX_NODE_TRAIN_FWD=999等,用于提升通信与执行效率。
六、性能基准(Benchmark)
1. Glint360K 的 MXNet 训练吞吐
在 8 × Tesla V100-SXM2-32GB 上训练 Glint360K 的吞吐量(img/sec):
| Backbone | GPU | FP16 | BatchSize/it | Throughput img/sec |
|---|---|---|---|---|
| R100 | 8 × V100-SXM2-32GB | False | 64 | 1748 |
| R100 | 8 × V100-SXM2-32GB | True | 64 | 3357 |
| R100 | 8 × V100-SXM2-32GB | False | 128 | 1847 |
| R100 | 8 × V100-SXM2-32GB | True | 128 | 3867 |
| R50 | 8 × V100-SXM2-32GB | False | 64 | 2921 |
| R50 | 8 × V100-SXM2-32GB | True | 64 | 5428 |
| R50 | 8 × V100-SXM2-32GB | False | 128 | 3045 |
| R50 | 8 × V100-SXM2-32GB | True | 128 | 6112 |
可以看出 FP16 混合精度可将吞吐提升约1.9~2 倍,且不损失精度——这也是文档强调"无性能下降"的重要实践支撑。
2. 百万级身份对比(忽略 IO 影响,混合精度训练,backbone=ResNet50)
100 万身份 / 8 × RTX2080Ti:
| Method | GPUs | BatchSize | Memory/M | Throughput img/sec | W |
|---|---|---|---|---|---|
| Model Parallel | 8 | 1024 | 10408 | 2390 | GPU |
| Partial FC (Ours) | 8 | 1024 | 8100 | 2780 | GPU |
1000 万身份 / 64 × RTX2080Ti:
| Method | GPUs | BatchSize | Memory/M | Throughput img/sec | W |
|---|---|---|---|---|---|
| Model Parallel | 64 | 2048 | 9684 | 4483 | GPU |
| Partial FC (Ours) | 64 | 4096 | 6722 | 12600 | GPU |
在千万身份规模下,Partial FC 相比模型并行将 batch size 从 2048 提升到 4096、显存占用从 9684M 降到 6722M、吞吐从 4483 提升到12600 img/sec(约 2.8 倍),充分体现了采样 + 分片策略在通信与显存上的双重优势。
七、训练速度问题排查(Troubleshooting)
官方文档给出了训练中常见的四类问题:
- Horovod 是否安装成功:执行
horovodrun --check,输出应显示[X] MXNet、[X] MPI、[X] Gloo、[X] NCCL等(完整输出示例见 mxnet/README.md); - MXNet 版本:部分版本与 horovod 存在兼容 Bug,建议使用 1.5 或 1.6;mxnet 1.5.1 无法安装 horovod;
- CUDA 版本一致性:mxnet-cu101 要求 CUDA 10.1,用
/usr/local/cuda/bin/nvcc -V核对(示例输出显示release 10.1, V10.1.168); - IO 瓶颈:
- 在
config中开启debug模式(config.debug = 1,对应--dataset 100w/1000w等调试配置),训练会改用DummyIter屏蔽真实数据读取,据此判断慢训练是否由 IO 引起; - 若确为 IO 瓶颈,可将数据集挂载到内存盘(tmpfs),例如 256G 内存的机器:
- 在
sudo mkdir /train_tmp mount -t tmpfs -o size=140G tmpfs /train_tmp八、PyTorch 版本与生态衔接
Partial FC 的 PyTorch 实现已合并进 arcface_torch(见 pytorch/README.md),相关训练代码、partial_fc_v2.py、损失函数与数据集加载均在 recognition/arcface_torch 中维护,同时仓库还提供配套的 torch2onnx.py、onnx_helper.py 与 eval_ijbc.py 用于模型导出与大规模评测,方便从训练到部署的完整衔接。
九、引用(Citation)
如果在研究中使用 Partial-FC 或 Glint360K,请引用以下论文:
@inproceedings{an_2022_pfc_cvpr, title={Killing Two Birds with One Stone: Efficient and Robust Training of Face Recognition CNNs by Partial FC}, author={An, Xiang and Deng, Jiangkang and Guo, Jia and Feng, Ziyong and Zhu, Xuhan and Jing, Yang and Tongliang, Liu}, booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition}, year={2022} } @inproceedings{an_2021_pfc_iccvw, title={Partial FC: Training 10 Million Identities on a Single Machine}, author={An, Xiang and Zhu, Xuhan and Gao, Yuan and Xiao, Yang and Zhao, Yongle and Feng, Ziyong and Wu, Lan and Qin, Bin and Zhang, Ming and Zhang, Debing and Fu, Ying}, booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) Workshops}, year={2021}, }总结
Partial FC 通过"分类层按 GPU 分片 + 负类中心采样"两板斧,将千万级身份的人脸识别训练从多机多卡的模型并行方案中解放出来:单机 8 卡即可完成 1000 万身份的 Softmax 分类训练,显存与通信开销双双下降、吞吐成倍提升。配合官方清洗发布的 Glint360K 数据集(1709 万图 / 36 万身份),在 IFRT、IJB-C、MegaFace 等评测上均取得了优于 MS1M-V2/V3 的成绩。本文所述的所有训练脚本、配置与实现均可直接在仓库 recognition/partial_fc 目录下复现与扩展。
【免费下载链接】insightfaceState-of-the-art 2D and 3D Face Analysis Project项目地址: https://gitcode.com/GitHub_Trending/in/insightface
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考