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
conda create -n python310 python=3.10
# 激活环境
conda activate python310
训练环境搭建和数据处理需要用到一些工具，安装命令如下:
sudo apt update
sudo apt install build-essential
sudo apt install portaudio19-dev
如果安装很慢可以替换为aliyun的mirror
sudo sed -i 's/http:\/\/archive.ubuntu.com/https:\/\/mirrors.aliyun.com/g' /etc/apt/sources.list
sudo apt update
训练环境
训练使用RAD-NeRF原作者kiui的开源代码,可以按照以下步骤进行：
- clone代码到本地
git clone https://github.com/ashawkey/RAD-NeRF.git --depth=1
cd RAD-NeRF
修改代码，data_utils/process.py 的 第 50 行，修改为：
fa = face_alignment.FaceAlignment(face_alignment.LandmarksType.TWO_D, flip_input=False)
原因是三方模块的字段变化.
- 安装依赖模块
# 设置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 scripts/install_ext.sh
也可以通过进入extension的代码目录进行编译安装，以raymarching这个扩展为例
cd raymarching
python setup.py --install
一共有4个扩展需要安装：freqencoder,gridencoder,shencoder,raymarching.
- 安装
pytorch3d
pip install "git+https://github.com/facebookresearch/pytorch3d.git"
注意编译安装耗时较长
- 下载数据处理所需的三方data
# 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/目录下. 并且执行
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行：
# 修改为自己的本地模型路径
self.processor = AutoProcessor.from_pretrained('/model/wav2vec')
self.model = AutoModelForCTC.from_pretrained('/model/wav2vec').to(self.device)
另一种方式是使用huggingface的镜像站，方法可以是在环境变量中指定HF_ENDPOINT
export HF_ENDPOINT=https://hf-mirror.com
其他方式参考镜像站的文档.
预处理
训练开始前需要对视频进行预处理，处理程序会将视频按帧切分，并且进行人脸区域的分离
首先将视频文件放到data目录下，例如视频名称为obama.mp4，就放在 data/obama/obama.mp4,然后执行：
python data_utils/process.py data/obama/obama.mp4
执行成功后可以看到目录结构为
|-- 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
python main.py data/obama/ --workspace trial_obama/ -O --iters 200000
注意这里的参数
data/obama 是数据目录,根据自己的文件命名
--workspace参数 trial_obama 是我们自己可以指定的值，能表达意思即可
- finetune lips
python main.py data/obama/ --workspace trial_obama/ -O --iters 250000 --finetune_lips
这里的路径要和第1步的保持一致，因为是在第1步的基础上进行的唇形 finetune.
- 训练torso
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.jsondata/bc.jpg-
trial_obama_torso/checkpoints/ngp_epxxxx.pth
这些文件根据需要放到执行推理所要求的路径下即可
