MotionBERT 源码复现

KangXi_TangYuan

已于 2023-12-05 22:05:45 修改

阅读量1.2k

点赞数 27

文章标签： python 人工智能计算机视觉

于 2023-12-05 22:03:37 首次发布

本文链接：https://blog.csdn.net/KangXi_TangYuan/article/details/134194719

版权

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 报错，需要下载这个包，随便解压到哪里都行，根据指引安装：
```
  python install -r requirements.txt
  python setup.py develop
```
不过这里它把我的torch删了下成了别的版本，回头可以手动改回来。
- 报错3：No such file or directory: 'data/AMASS/body_models/smplh/female/model.npz'，下载模型解压到 AMASS/body_models/smplh。
- 报错4：No such file or directory: 'data/AMASS/body_models/dmpls/female/model.npz'，下载dmpl解压到AMASS/body_models/dmpls
- 以上，可以成功运行preprocess_amass.py了，大约耗时2h。
最后运行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：
```
#subset_list: [AMASS, H36M-SH]
train_subset_list: [AMASS, H36M-SH]
test_subset_list: [H36M-SH]
```
并在train.py里的train_dataset和test_dataset相应位置改一下：
```
# 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