MotionBERT复现流程
源码地址:https://github.com/Walter0807/MotionBERT/blob/main/docs/pretrain.md
一、配置环境、数据集
1. 根据指引安装环境
安装虚拟环境
安装pytorch时不用conda,用pip install
pip install torch==1.13.0+cu116 torchvision==0.14.0+cu116 torchaudio==0.13.0 --extra-index-url https://download.pytorch.org/whl/cu116
2. 下载数据集
在MotionBERT-main下创建data文件夹。
(1) H36M:
将下载的h36m的zip包解压到data/motion3d(一个.kpl文件)。运行python tools/convert_h36m.py
。
报错:No module named ‘yaml’
解决:pip install pyyaml
(2) PoseTrack18:
下载annotation版本,解压至data/motion2d/posetrack8_annotations
(3) InstaVariety:
直接解压到data/motion2d
(4) AMASS:
前三个数据集作者都提供了处理好的数据,但AMASS没有,按照作者给的流程到官网下载数据集并预处理。
- 先下载所有含SMPL-H的数据集(GRAB和SOMA除外,这两个的确没有dmpls参数),挨个解压到data/AMASS/amass_202203文件夹下(这是作者起的名字,也可以改掉),总计20个文件夹(删掉多余license)。
- 运行compress_amass.py,会在AMASS文件夹下生成一些.pkl和.csv文件。
在这里发现2080处理数据比3090快很多。
运行的时候经常无缘无故killed,应该是内存满了(30G都不够),多申一些(60G)就够了
- 运行preprocess_amass.py会有一些报错,需要下载一些文件。
- 报错1:将J_regressor_h36m_correct.npy放入AMASS文件夹下。
- 报错2:导包时
from human_body_prior.body_model.body_model import BodyModel
报错,需要下载这个包,随便解压到哪里都行,根据指引安装:
不过这里它把我的torch删了下成了别的版本,回头可以手动改回来。python install -r requirements.txt python setup.py develop
- 最后运行convert_amass.py。会在motion3d/MB3D_f243s81下生成AMASS/train/…(96179个.pkl文件)
(5) Mesh数据集
根据指导下载数据集和smpl模型并放入data/mesh文件夹下。
将torch, torchvision, torchaudio改回原来的版本。
二、Running
1. 修改train/test数据集
- 对于pretrain,直接运行会发现motion3d下的AMASS里只有train,没有test导致报错。但作者只用了h36m进行了test,AMASS里就是没有test。所以稍微改一下yaml文件中的subset_list:
并在train.py里的train_dataset和test_dataset相应位置改一下:#subset_list: [AMASS, H36M-SH] train_subset_list: [AMASS, H36M-SH] test_subset_list: [H36M-SH]
# 237行 train_dataset = MotionDataset3D(args, args.train_subset_list, 'train') test_dataset = MotionDataset3D(args, args.test_subset_list, 'test')
- 同理修改configs/pose3d/里的.yaml文件:
#subset_list: [H36M-SH] train_subset_list: [H36M-SH] test_subset_list: [H36M-SH]
- 发现3090读取数据直接卡死,用2080只需要10秒,修改batch_size。
- 发现for循环里tqdm(enumerate(data_loader))不显示进度条,将两者调换一下顺序即可:enumerate(tqdm(data_loader))
2. 添加wandb
先在wandb中添加project:MotionBERT
在train.py文件中添加wandb:
import wandb
if __name__ == "__main__":
wandb.login()
wandb.init(project='MotionBERT', entity="...", name='...')
...
寻找training loss,发现在train_epoch()中,在相应的3d和2d loss后添加:
wandb.log({
'loss_3d_pos': loss_3d_pos.item(),
'loss_3d_scale': loss_3d_scale.item(),
'loss_3d_velocity': loss_3d_velocity.item(),
'loss_lv': loss_lv.item(),
'loss_lg': loss_lg.item(),
'loss_a': loss_a.item(),
'loss_av': loss_av.item(),
'loss_3d_total': loss_total.item()
})
wandb.log({
'loss_2d_proj': loss_2d_proj.item(),
'loss_2d_total': loss_total.item()
第一次登录需要API,40个字符长,但怎么复制粘贴都会有问题,解决方法是直接手打(全程不会显示)
train_mesh.py文件同理
3. 下载模型
训练好的模型作者提供了pretrain和pose3d,mesh只提供了一个ft_pw3d,下载好放到相应位置。
因此mesh中的finetune只能跑:
# with 3DPW
python train_mesh.py \
--config configs/mesh/MB_ft_pw3d.yaml \
--pretrained checkpoint/pretrain/MB_release \
--checkpoint checkpoint/mesh/FT_MB_release_MB_ft_pw3d
evaluate只能跑:
# with 3DPW
python train_mesh.py \
--config configs/mesh/MB_ft_pw3d.yaml \
--evaluate checkpoint/mesh/FT_MB_release_MB_ft_pw3d/best_epoch.bin
其他需要自己训练
4. 解决killed问题
运行train_mesh.py时在进行到evaluate_mesh()步骤时程序总是killed,进入到lib/utils/utils_mesh.py下debug,发现在运行#mpve = np.mean(np.mean(np.sqrt(np.square(pred_verts - target_verts).sum(axis=2)), axis=1))
时kill,具体原因是pred_verts和
target_verts都是非常大的矩阵,再进行square操作会超内存。让chatgpt帮忙改一下,改成分块计算的方式:
#mpve = np.mean(np.mean(np.sqrt(np.square(pred_verts - target_verts).sum(axis=2)), axis=1))
chunk_size = 1000
result = compute_squared_difference(pred_verts, target_verts, chunk_size)
mpve = np.mean(np.mean(np.sqrt(result), axis=1))
compute_squared_difference()的定义:
def compute_squared_difference(pred_verts, target_verts, chunk_size):
num_samples, num_verts, num_dims = pred_verts.shape
result = np.zeros((num_samples, num_verts), dtype=np.float32)
for i in range(0, num_samples, chunk_size):
start = i
end = min(i + chunk_size, num_samples)
chunk_pred = pred_verts[start:end, :, :]
chunk_target = target_verts[start:end, :, :]
chunk_squared_diff = np.sum(np.square(chunk_pred - chunk_target), axis=2)
result[start:end, :] = chunk_squared_diff
return result
此外,申请的内存改为80G就可以正常使用了(之前申请了60G竟然还不够)
5. 附所有运行过程
至此,所有的范例代码都能跑通了。所有代码都先用2080确认无误,需要的话再后续移植到A100。
<1> pretrain
(1) MB_pretrain: (有)
python train.py --config configs/pretrain/MB_pretrain.yaml -c checkpoint/pretrain/MB_pretrain
2080: bs=2
<2> pose3d
(1) Train from scratch: (有)
python train.py --config configs/pose3d/MB_train_h36m.yaml --checkpoint checkpoint/pose3d/MB_train_h36m
(2) Finetune: (有)
python train.py --config configs/pose3d/MB_ft_h36m.yaml --pretrained checkpoint/pretrain/MB_release --checkpoint checkpoint/pose3d/FT_MB_release_MB_ft_h36m
(3)Evaluate:
# train (有)
python train.py --config configs/pose3d/MB_train_h36m.yaml --evaluate checkpoint/pose3d/MB_train_h36m/best_epoch.bin
# ft (有)
python train.py --config configs/pose3d/MB_ft_h36m.yaml --evaluate checkpoint/pose3d/FT_MB_release_MB_ft_h36m/best_epoch.bin
<3> mesh
(1) Train from scratch:
# with 3DPW (无)
python train_mesh.py --config configs/mesh/MB_train_pw3d.yaml --checkpoint checkpoint/mesh/MB_train_pw3d
# H36M (无)
python train_mesh.py --config configs/mesh/MB_train_h36m.yaml --checkpoint checkpoint/mesh/MB_train_h36m
(2) Finetune:
# with 3DPW (有)
python train_mesh.py --config configs/mesh/MB_ft_pw3d.yaml --pretrained checkpoint/pretrain/MB_release --checkpoint checkpoint/mesh/FT_MB_release_MB_ft_pw3d
# H36M (无)
python train_mesh.py --config configs/mesh/MB_ft_h36m.yaml --pretrained checkpoint/pretrain/MB_release --checkpoint checkpoint/mesh/FT_MB_release_MB_ft_h36m
(3)Evaluate:
# with 3DPW -train
python train_mesh.py --config configs/mesh/MB_train_pw3d.yaml --evaluate checkpoint/mesh/MB_train_pw3d/best_epoch.bin
# H36M -train
python train_mesh.py --config configs/mesh/MB_train_h36m.yaml --evaluate checkpoint/mesh/MB_train_h36m/best_epoch.bin
# with 3DPW -ft
python train_mesh.py --config configs/mesh/MB_ft_pw3d.yaml --evaluate checkpoint/mesh/FT_MB_release_MB_ft_pw3d/best_epoch.bin
# H36M -ft
python train_mesh.py --config configs/mesh/MB_ft_h36m.yaml --evaluate checkpoint/mesh/FT_MB_release_MB_ft_h36m/best_epoch.bin