使用Jittor框架实现的神经网络模型，能够从3D点云数据（顶点和法线（可选））预测关节位置和蒙皮权重。经过 10000 epochs 预训练与 2240 epochs 微调，可以取得 j2j_loss：0.0264, skin_l1loss：0.0046, vertex_loss：0.0498 , score：51.8823。我们后续微调到了 10000 epochs，但是没有时间进行测试了。
本仓库包含了一个神经网络模型的实现，该模型以3D点云数据（顶点和法线（可选））作为输入，并预测：
-
初始关节位置
-
关节位置细化偏移量（deltas）
-
每个顶点的蒙皮权重
-
傅里叶位置编码，提供丰富的位置信息
-
多层级特征提取，包含 Down Sample & Local Pooling 操作
-
自注意力块（Self Attention + FFN）增强特征表示
-
输出多尺度特征用于解码
- 使用简易 MLP 进行初始关节预测
- 初始关节同样会被
loss监督
- 使用
detach()后的初始关节作为输入 - Self Attention + FFN 建模关节偏移量之间的关系
- Cross Attention + FFN 根据点云特征调整关节偏移量
-
类 Unet 架构，逐步插值的同时融合多尺度特征
-
带自注意力的解码块
项目基本继承baseline的结构，具体如下：
jittor-comp-human/
├── data/
├── dataset/
│ ├── asset.py 加入了对点云整体进行随机旋转的方法
│ ├── dataset.py 加入了动捕动作&随机旋转的数据增强方法
│ ├── exporter.py
│ ├── format.py
│ └── sampler.py 使用了新的sample方式
├── launch/
│ ├── pretrain_joint_model.sh 预训练启动脚本
│ ├── finetune_joint_model.sh 微调启动脚本
│ └── predict_joint_model.sh 推理启动脚本
├── models/
│ ├── joint_model.py 主模型文件
│ ├── FourierEmbedder.py Fourier Embedding层
│ ├── transformer_block.py 使用的Self-Attention与Cross-Attention模块
│ └── metrics.py
├── output/ ckpt保存点
├── PCT/
├── README.md
├── predict_joint_model.py
└── train_joint_model.py本项目预训练阶段在 2 张 4090 上运行，训练时间约为 60 小时
由于时间不足，微调阶段使用了 8 张 A100，时间约为 11 小时
- Ubuntu 20.04.3 LTS
- CUDA 12.2
- python 3.11.13
- jittor 1.3.10.0
执行以下命令安装依赖
conda create -f environment.yml
预训练模型下载地址为 https://cloud.tsinghua.edu.cn/d/e6bd32bcec0441178ba2/ ，下载后放入目录 <root>/output/TransRefiner-finetune/
将数据下载解压到 <root>/data 下即可，下载链接：https://cloud.tsinghua.edu.cn/f/a0e48edfe8834c7b8b4c/?dl=1
我们首先进行了 10000 epochs 的预训练，这一阶段对 50% 的数据进行随机姿态的数据增强。
单卡可运行：
bash launch/pretrain_joint_model.sh
多卡可运行（这里以 2 张卡举例）：
mpirun -np 2 bash ./launch/pretrain_joint_model.sh
预训练后，我们针对提供的动捕动作进行了 2240 epochs 的微调（为什么是这个数呢，因为后来没时间了只训到了这里）。这一阶段的数据增强是 50% 的动捕动作 + 35% 的随机动作 + 15% 的原始姿态。此外，以 p=0.2 对整个点云进行随机旋转。
与预训练类似，把脚本换成 ./launch/finetune_joint_model.sh 即可
将模型放到 <root>/output/TransRefiner-finetune/ 下后，执行
bash ./launch/predict_joint_model.sh
生成的结果会放在 <root>/predict_joint/Transrefiner_resume_12240 下