RAD-NeRF数字人模型训练详解

RAD-NeRF是实现数字人的一个非常优秀的模型，在发布时达到SOTA.

本文记录RAD-NeRF模型训练的详细过程，包括系统环境介绍，训练环境搭建及完整训练流程.

硬件要求

训练需要使用CUDA, GPU 显存至少需要24G,训练中不同阶段占用显存不同，峰值超过22G.

系统环境

本文使用Ubuntu 22.04, Python 3.10, Pytorch 2.0.1, CUDA 11.7 作为训练的基础环境，并且在docker容器中亦进行了测试确认可工作。以下讨论的各种工具安装皆默认为Ubuntu 22.04.

如果需要构建训练容器镜像，可以使用nvidia发布的nvidia/cuda:11.7.1-devel-ubuntu22.04 作为基础镜像.

为了创建一个干净的python环境预防模块版本导致的问题，推荐使用Anaconda, 安装可以参考Installation -- Anaconda Doc

安装后执行如下命令创建环境并指定python版本为3.10

ini 复制代码

 conda create -n python310 python=3.10
 # 激活环境
 conda activate python310

训练环境搭建和数据处理需要用到一些工具，安装命令如下:

sql 复制代码

 sudo apt update
 sudo apt install build-essential
 sudo apt install portaudio19-dev

如果安装很慢可以替换为aliyun的mirror

bash 复制代码

 sudo sed -i 's/http://archive.ubuntu.com/https://mirrors.aliyun.com/g' /etc/apt/sources.list
 sudo apt update

训练环境

训练使用RAD-NeRF原作者kiui的开源代码,可以按照以下步骤进行：

clone代码到本地

bash 复制代码

 git clone https://github.com/ashawkey/RAD-NeRF.git --depth=1
 cd RAD-NeRF

修改代码，data_utils/process.py 的第 50 行，修改为：

ini 复制代码

 fa = face_alignment.FaceAlignment(face_alignment.LandmarksType.TWO_D, flip_input=False)

原因是三方模块的字段变化.

安装依赖模块

csharp 复制代码

 # 设置pip使用aliyun mirror
 pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/
 # 提前安装pytorch指定为2.0.1
 pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2
 # 安装 requirements.txt中的依赖
 pip install -r requirements.txt

编译代码库中的python extension

bash 复制代码

 bash scripts/install_ext.sh

也可以通过进入extension的代码目录进行编译安装，以raymarching这个扩展为例

arduino 复制代码

 cd raymarching
 python setup.py --install

一共有4个扩展需要安装：freqencoder,gridencoder,shencoder,raymarching.

安装pytorch3d

arduino 复制代码

 pip install "git+https://github.com/facebookresearch/pytorch3d.git"

注意编译安装耗时较长

下载数据处理所需的三方data

bash 复制代码

 # face parsing
 wget https://github.com/YudongGuo/AD-NeRF/blob/master/data_util/face_parsing/79999_iter.pth?raw=true -O data_utils/face_parsing/79999_iter.pth
 # AD-NeRF
 wget https://github.com/YudongGuo/AD-NeRF/blob/master/data_util/face_tracking/3DMM/exp_info.npy?raw=true -O data_utils/face_tracking/3DMM/exp_info.npy
 wget https://github.com/YudongGuo/AD-NeRF/blob/master/data_util/face_tracking/3DMM/keys_info.npy?raw=true -O data_utils/face_tracking/3DMM/keys_info.npy
 wget https://github.com/YudongGuo/AD-NeRF/blob/master/data_util/face_tracking/3DMM/sub_mesh.obj?raw=true -O data_utils/face_tracking/3DMM/sub_mesh.obj
 wget https://github.com/YudongGuo/AD-NeRF/blob/master/data_util/face_tracking/3DMM/topology_info.npy?raw=true -O data_utils/face_tracking/3DMM/topology_info.npy

访问 basel face model 下载这个链接，填写信息后下载01_MorphableModel.mat,放置到data_utils/face_tracking/3DMM/目录下. 并且执行

bash 复制代码

 cd data_utils/face_tracking
 python convert_BFM.py

至此训练环境搭建完成

模型训练

素材准备

训练要求素材为 fps 25, 分辨率 512*512 的视频，时长 3~5 min. 经过实测，对分辨率的要求并不严格，测试使用例如720*720的分辨率也可以成功.

ASR 模型准备

笔者使用 wav2vec作为ASR模型，代码中直接从huggingface加载模型，如果因为网络问题无法直接使用，也可以先下载模型到本地，再修改代码为从本地路径下载，修改代码处为nerf/asr.py 92~93行：

ini 复制代码

 # 修改为自己的本地模型路径
 self.processor = AutoProcessor.from_pretrained('/model/wav2vec')
 self.model = AutoModelForCTC.from_pretrained('/model/wav2vec').to(self.device)

另一种方式是使用huggingface的镜像站，方法可以是在环境变量中指定HF_ENDPOINT

arduino 复制代码

 export HF_ENDPOINT=https://hf-mirror.com

其他方式参考镜像站的文档.

预处理

训练开始前需要对视频进行预处理，处理程序会将视频按帧切分，并且进行人脸区域的分离首先将视频文件放到data目录下，例如视频名称为obama.mp4，就放在 data/obama/obama.mp4,然后执行：

bash 复制代码

 python data_utils/process.py data/obama/obama.mp4

执行成功后可以看到目录结构为

lua 复制代码

 |-- aud.wav
 |-- aud_eo.npy
 |-- bc.jpg
 |-- obama.mp4
 |-- gt_imgs
 |-- ori_imgs
 |-- parsing
 |-- torso_imgs
 |-- track_params.pt
 |-- transforms_train.json
 `-- transforms_val.json

训练

训练分为3步

训练head

css 复制代码

 python main.py data/obama/ --workspace trial_obama/ -O --iters 200000

注意这里的参数 data/obama 是数据目录,根据自己的文件命名 --workspace参数 trial_obama 是我们自己可以指定的值，能表达意思即可

finetune lips

css 复制代码

 python main.py data/obama/ --workspace trial_obama/ -O --iters 250000 --finetune_lips

这里的路径要和第1步的保持一致，因为是在第1步的基础上进行的唇形 finetune.

训练torso

css 复制代码

 python main.py data/obama/ --workspace trial_obama_torso/ -O --torso --head_ckpt <head>.pth --iters 200000

这里的--worksapce参数trial_obama_torso 是自行指定的一个新目录，用来放torso的权重文件，--head_ckpt 值是第2步的产出物，在我们例子中，就是 trial_obama/checkpoints/ngp_ep0084.pth，数字可能会不同，因为视频不同，经过的epoch也会不同，文件中的数字代表了训练的epoch.

模型文件

训练完成后，推理阶段需要的文件有：

data/transforms_train.json
data/bc.jpg
trial_obama_torso/checkpoints/ngp_epxxxx.pth 这些文件根据需要放到执行推理所要求的路径下即可

参考文档

github.com/ashawkey/RA...