Jittor_骨骼绑定赛题_undefined队

使用Jittor框架实现的神经网络模型，能够从3D点云数据（顶点和法线（可选））预测关节位置和蒙皮权重。经过 10000 epochs 预训练与 2240 epochs 微调，可以取得 j2j_loss：0.0264, skin_l1loss：0.0046, vertex_loss：0.0498 , score：51.8823。我们后续微调到了 10000 epochs，但是没有时间进行测试了。

alt text

简介

本仓库包含了一个神经网络模型的实现，该模型以3D点云数据（顶点和法线（可选））作为输入，并预测：

初始关节位置
关节位置细化偏移量（deltas）
每个顶点的蒙皮权重

模型架构

alt text

主干网络（PT2_Backbone）

傅里叶位置编码，提供丰富的位置信息
多层级特征提取，包含 Down Sample & Local Pooling 操作
自注意力块（Self Attention + FFN）增强特征表示
输出多尺度特征用于解码

关节回归（Joint Regression）

使用简易 MLP 进行初始关节预测
初始关节同样会被 loss 监督
关节细化层（Joint Refinement Layer）
使用 detach() 后的初始关节作为输入
Self Attention + FFN 建模关节偏移量之间的关系
Cross Attention + FFN 根据点云特征调整关节偏移量

蒙皮权重解码器（Skin Decoder）

类 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 下