InsightFace Partial FC 大规模人脸识别分布式训练实战指南MXNet 版【免费下载链接】insightfaceState-of-the-art 2D and 3D Face Analysis Project项目地址: https://gitcode.com/GitHub_Trending/in/insightface本文以 InsightFace 仓库中 recognition/partial_fc/mxnet 的官方训练文档为主线完整讲解基于 Partial FCPartial Fully Connected Layer方法的大规模人脸识别分布式训练从环境安装、版本兼容、多机多卡启动到分类层模型并行与 Softmax 近似采样的底层原理再到性能瓶颈排查与内存盘加速。读完本文你将能独立搭建环境并复现百万到千万级类别class身份分类任务的 MXNet Horovod 训练流程并理解其8 卡 RTX 2080Ti 即可训练千万级类别的核心机制。Partial FC 论文对应 arXiv 2203.15565其核心思路是当人脸训练集的类别数从百万级增长到亿级时传统 Softmax 分类层权重矩阵与全类别计算会成为显存和通信瓶颈Partial FC 通过分类层模型并行 正类必选、负类随机采样的 Softmax 近似将训练成本从与类别总数线性相关降低为与采样数相关实现无损loss-free的大规模训练。一、Partial FC 核心原理为什么能训练亿级类别在了解如何安装运行之前先理解这套训练框架的设计动机。MXNet 版本 Partial FC 的实现代码全部位于 recognition/partial_fc/mxnet 目录下核心包括memory_bank.py内存库Memory Bank负责在每张卡上均匀存储一部分类中心及其动量memory_softmax.py带 margin 的 Softmax 损失ArcFace / CosFacememory_module.py分布式采样训练模块SampleDistributeModule承担特征 AllGather、损失与梯度同步memory_samplers.py正类必选 负类随机采样器train_memory.py训练入口脚本。1. 分类层模型并行全程只需三次通信官方文档英文版 README.md 的 Our Method 一节明确指出类中心被均匀分布到不同 GPU上完成一次无损失的 Softmax 计算只需要三次通信特征同步AllGather保证每张 GPU 都拥有全局所有 GPU 上的特征。对应源码中 memory_module.py 的allgather实现——该实现利用 in-place 的hvd.allreduce_(total_tensor, averageFalse)完成 AllGather先把本卡特征写入total_tensor中本 rank 对应的分片再通过一次全归约让所有卡都拿到完整拼接后的全局特征。Softmax 分母同步Allreduce先在各卡本地计算exp(logits)的部分和再通过一次通信求出全局和。对应 memory_module.py 中先对每个 logit 减去全局最大值数值稳定再hvd.allreduce(sum_fc7, averageFalse)汇总全局分母。特征梯度同步Allreducelogits 的梯度可以独立计算特征梯度同理最后把各卡梯度收集回来送回 backbone。对应 memory_module.pytotal_feature_grad hvd.allreduce(total_feature_grad, averageFalse)后每个 rank 切出属于自己的分片fc1_grad传给 backbone 做反向传播。在训练循环中每次迭代都会执行forward_backward前向 反向与update更新 backbone 参数这一调用链在 memory_module.py 的fit方法中实现。2. Softmax 近似正类必选、负类随机采样当类别数极大时参与 Softmax 计算的类中心子集即可近似完整 Softmax正类中心必须包含在子集中。官方文档给出的伪代码如下centers_p func_positive(label) # 根据样本标签选出正类中心 centers_n func_negative(centers_p) # 排除正类后随机采样负类中心 centers_final concat(centers_n, centers_p) # 参与 Softmax 计算的类中心这一逻辑在 memory_samplers.py 中被拆成三个类精确实现CenterPositiveClassGet从全局标签中筛选出落在本 rank 类中心区间[rank * num_local, rank * num_local num_local)内的正类标签CenterNegetiveClassSample在排除正类后从本卡num_local个类中心的池子里用np.random.choice(..., replaceFalse)无放回采样凑满num_sample个WeightIndexSampler将正类索引与负类索引拼接并断言最终数量、去重数量均等于num_sample保证每个 batch 参与计算的类中心数严格受控。是否启用采样由config.sample_ratio控制当sample_ratio 1.0时走完整 Softmax 的backward分支否则走采样分支backward_sample见 memory_module.py。二、环境安装从零搭建训练环境官方中文文档 README_CN.md 将安装拆解为 5 个步骤建议按顺序执行。1. Python 依赖pip install easydict mxboard opencv-python tqdm各依赖的作用easydict以属性方式访问配置字典训练配置 default.py 即基于EasyDict定义mxboard训练日志可视化回调 callbacks.py 中用于写 Summaryopencv-python数据增强与图像读取FaceImageIter依赖tqdm进度条显示。2. 安装 NCCL可选但推荐NCCL 不装也能跑但装上后 GPU 通信更快。NCCL 安装需要与 CUDA 版本对应。config.sh中通过HOROVOD_GPU_ALLREDUCENCCL、HOROVOD_GPU_ALLGATHERNCCL显式启用 NCCL 后端请参考 NVIDIA 官方 NCCL 安装指南选择对应版本。3. 安装 OpenMPI必须OpenMPI 必须安装官方强烈建议使用仓库自带脚本 setup-utils/install-mpi.sh 编译源码安装OpenMPI 4.0.0不要用系统包管理器直接安装否则可能与 Horovod 链接的 MPI 不一致导致运行时异常。4. 安装 Horovod 与 MXNet版本兼容性最关键官方文档明确警告某些版本的 MXNet 无法安装 Horovod并给出了严格的版本兼容表mxnethorovodcuda1.4.0xx1.5.0可以安装cuda10.01.5.1xx1.6.0.post0可以安装cuda10.11.7.0xx强烈建议使用mxnet1.6.0对应mxnet-cu1011.6.0.post0与cuda10.1。英文版 README 同时补充社区发现 mxnet 1.5.1 无法安装 Horovod。Horovod 的安装同样推荐使用仓库自带脚本 setup-utils/install-horovod.sh对应版本 Horovod 0.19.2。安装完成后运行以下命令检查是否成功NCCL 可选有则更快# Horovod v0.19.2: # Available Frameworks: # [ ] TensorFlow # [ ] PyTorch # [X] MXNet # # Available Controllers: # [X] MPI # [X] Gloo # # Available Tensor Operations: # [X] NCCL # [ ] DDL # [ ] CCL # [X] MPI # [X] Gloo其中 MXNet 框架已勾选、MPI/Gloo 控制器可用、NCCL 张量操作可用即表示安装成功。5. SSH 无密登录多机分布式训练时每台机器都需要与所有机器包括自己建立免密登录。官方推荐使用一条命令完成ssh-copy-id userip三、训练配置解析default.py 与启动脚本1. 训练入口与命令行参数训练入口 train_memory.py 支持 4 个命令行参数参数默认值说明--datasetemore数据集配置如emore、glint360k_8GPU、glint360k_16GPU--networkr100骨干网络如r50、r100、r122、rx50、rx101--losscosface损失配置arcface或cosface--models-root./test模型保存根目录2. 核心超参数default.pydefault.py 中的全局配置配置项默认值说明embedding_size512特征维度image_size112输入图像尺寸112×112batch_size64每卡 batch sizebackbone_lr0.1骨干网络学习率memory_bank_lr等于 backbone_lr类中心学习率sample_ratio1.0负类采样比例1.0时启用近似 Softmaxdebug0开启后使用DummyIter屏蔽真实 IOfp16False是否混合精度训练loss_s/loss_m1/loss_m2/loss_m3依 loss 而定margin 损失参数ArcFaces64, m11.0, m20.5, m30.0CosFaces64, m11.0, m20.0, m30.4数据与训练步数按数据集分组配置例如emoreMS1MV385742 类lr_steps100000,160000max_update180000batch_size64glint360k_8GPU360232 类lr_steps200000,400000,500000,550000max_update600000batch_size64验证集为agedb_30, calfw, cfp_ff, cfp_fp, cplfw, lfw, vgg2_fpglint360k_16GPUlr_steps200000,280000,360000max_update400000webface10575 类lr_steps20000,28000max_update32000100w / 1000w / 2000w / 3000w / 10000w用于大规模 Debug 的模拟配置debug1类别数分别为 100 万到 1 亿。训练中的学习率调度由 memory_scheduler.py 实现基于config.lr_steps生成两个MultiFactorSchedulerstep 处乘以factor0.1分别用于 backbone 与类中心。3. 启动脚本 config.shconfig.sh 是每进程实际执行的训练脚本内容如下export CUDA_VISIBLE_DEVICES0,1,2,3,4,5,6,7 export HOROVOD_GPU_ALLREDUCENCCL export HOROVOD_GPU_ALLGATHERNCCL export HOROVOD_GPU_BROADCASTNCLL export MXNET_CPU_WORKER_NTHREADS3 # use which python to get the absolute path of your python interpreter 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注意其中HOROVOD_GPU_BROADCASTNCLL为原脚本中的拼写实际生效的 NCCL 相关配置是前两行的ALLREDUCE与ALLGATHER保持原样即可无需改动。训练入口 train_memory.py 还会在进程启动时设置一批 MXNet / Horovod 环境变量例如关闭镜像反向传播MXNET_BACKWARD_DO_MIRROR0、关闭 KVStore 更新MXNET_UPDATE_ON_KVSTORE0、启用批量算子MXNET_EXEC_BULK_EXEC_MAX_NODE_TRAIN_FWD999、MXNET_EXEC_BULK_EXEC_MAX_NODE_TRAIN_BWD25、Horovod 融合阈值 64MBHOROVOD_FUSION_THRESHOLD67108864、2 个 NCCL 流等这些都是针对大规模分布式训练调优过的默认值。四、启动分布式训练horovodrun 与 mpirun 两种方式Horovod 底层调用的仍是 MPI有多少块 GPU就要启动多少个进程。官方提供两种启动方式。1. 使用 horovodrun 运行单机 8 卡horovodrun -np 8 -H localhost:8 bash config.sh两台机器共 16 卡horovodrun -np 16 -H ip1:8,ip2:8 bash config.sh其中-np指定进程总数-H指定主机与每台机器的进程数即 GPU 数。多机场景要求执行horovodrun的主机必须能免密 SSH 到所有其他主机且无任何交互提示。2. 使用 mpirun 运行仓库提供了封装好的 run.shbash run.sh其内容为mpirun -np 8 \ -hostfile hosts/host_8 --allow-run-as-root \ -bind-to none -map-by slot \ -x LD_LIBRARY_PATH -x PATH \ -mca pml ob1 -mca btl ^openib \ -mca btl_tcp_if_include eth0 \ -x OMP_NUM_THREADS2 \ bash config.sh关键参数说明-hostfile hosts/host_8主机列表文件仓库同时提供 hosts/host_8单机 8 卡与 hosts/host_16双机 16 卡两个模板多机场景请将文件内容替换为实际 IP-bind-to none -map-by slot进程不与具体 CPU 核绑定按槽位映射避免 MPI 进程间的资源争抢-x LD_LIBRARY_PATH -x PATH将环境变量透传给各进程-mca pml ob1 -mca btl ^openib使用 ob1 点对点消息层并禁用 openibInfiniBand传输-mca btl_tcp_if_include eth0限定 TCP 走 eth0 网卡多机场景按实际网卡名调整-x OMP_NUM_THREADS2限制 OpenMP 线程数避免与 GPU 计算抢 CPU。3. 训练中的验证与模型保存训练过程中local_rank 0的进程会通过CallBackVertification在config.verbose默认 2000 步间隔执行验证集测试CallBackModelSave保存 backbone 模型CallBackCenterSave把每张卡的类中心保存为{rank}_centers.param与{rank}_centers_mom.param见 callbacks.py 与 memory_bank.py。五、错误排查与性能调优1. 检查 Horovod 是否安装成功运行horovodrun --check输出如第三节所示MXNet 框架勾选、MPI/Gloo 控制器可用、NCCL 可用即为成功。2. 检查 CUDA 版本与 MXNet 是否匹配例如mxnet-cu101需要 CUDA 10.1/usr/local/cuda/bin/nvcc -V # nvcc: NVIDIA (R) Cuda compiler driver # Copyright (c) 2005-2019 NVIDIA Corporation # Built on Wed_Apr_24_19:10:27_PDT_2019 # Cuda compilation tools, release 10.1, V10.1.1683. 排查 IO 是否拖慢训练可以在config.py即 default.py中开启config.debug 1。开启后训练入口会使用DummyIter替代真实的FaceImageIter见 train_memory.py完全屏蔽数据读取从而判断训练速度瓶颈是否来自 IO。default.py中100w / 1000w / 2000w / 3000w / 10000w数据集配置即预设了debug1可直接用于 IO 压测。4. 数据挂载到内存盘加速如果确认瓶颈是 IO可以把数据集挂载到 tmpfs 内存盘前提是 RAM 足够大官方示例为 256G 内存挂 140G# 如果你的 RAM 有 256G sudo mkdir /train_tmp mount -t tmpfs -o size140G tmpfs /train_tmp挂载后把train.rec放到/train_tmp下并将 default.py 中各数据集的config.rec路径指向/train_tmp/.../train.recwebface、glint360k、emore配置中已预留了/train_tmp路径。5. 其他注意事项日志每个 rank 的训练日志写入models_root/{rank}_hist.log多卡排错时可分 rank 查看对应日志验证集数据验证集.bin文件需与train.rec位于同一目录os.path.dirname(config.rec)CallBackVertification按config.val_targets自动加载环境限制该套 MXNet 方案面向 2019 年前后软件栈CUDA 10.1 / mxnet-cu101 / Horovod 0.19.2 / OpenMPI 4.0.0新硬件与新版驱动环境下需注意版本适配。六、总结Partial FC 通过类中心按 rank 均匀分片存储于各 GPU 每次迭代仅三次通信 正类必选负类采样的 Softmax 近似三重设计将人脸识别分类层的显存与通信开销从类别总数维度降为采样数维度。本文所述的 MXNet 实现recognition/partial_fc/mxnet提供了从环境安装NCCL / OpenMPI / Horovod / MXNet 版本矩阵、配置解析default.py、双启动方式horovodrun/mpirun到故障排查horovodrun --check、CUDA 匹配、debug 模式、tmpfs 内存盘的完整闭环是复现大规模人脸识别训练或将其迁移到新框架时的重要参考实现。【免费下载链接】insightfaceState-of-the-art 2D and 3D Face Analysis Project项目地址: https://gitcode.com/GitHub_Trending/in/insightface创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考