3. release training code

zju3dv · Nov 2, 2024 · 5a67f9d · 5a67f9d
1 parent 742e81c
commit 5a67f9d
Showing 27 changed files with 1,612 additions and 5 deletions.
diff --git a/README.md b/README.md
@@ -5,9 +5,11 @@ This repo is the official implementation of "Generating Human Motion in 3D Scene
 ![pipeline](doc/pipeline.png)
 
 ## News
+[2024/11/02] We release the training code.
+
 [2024/10/21] We release the visualization code.
-[2024/06/09] We first release the test & evaluation code.
 
+[2024/06/09] We first release the test & evaluation code.
 ## Installation
 ```bash
 conda create -n most python=3.9
@@ -47,6 +49,15 @@ mkdir data
 ln -s /path/to/humanise data/HUMANISE
 ```
 
+### AMASS dataset
+(Only needed if you want to train the models by yourself.)
+1. Please follow [HUMOR](https://github.com/davrempe/humor/tree/main?tab=readme-ov-file#datasets) to download and preprocess AMASS dataset.
+2. Link to data/
+```bash
+ln -s /path/to/amass_processed data/amass_preprocess
+```
+
+
 ### SMPLX models
 1. Download SMPLX models from [link](https://smpl-x.is.tue.mpg.de/).
 2. Put the smplx folder under ```data/smpl_models``` folder:
@@ -91,7 +102,7 @@ The generated results are shared in [link](https://drive.google.com/file/d/1zrpz
 We use [wis3d](https://pypi.org/project/wis3d/) lib to visualize the results.
 To prepare for the visualization:
 ```bash
-python tools/visualizae_results.py -c configs/test/visualize.yaml
+python tools/visualize_results.py -c configs/test/visualize.yaml
 ```
 Then, in terminal:
 ```bash
@@ -100,6 +111,26 @@ wis3d --vis_dir out/vis3d --host ${HOST} --port ${PORT}
 You can then visualize the results in ```${HOST}:${PORT}```.
 
 
+# Train the models by yourself
+## Pretrain on the AMASS dataset
+Train the trajectory model:
+```bash
+python tools/train.net -c configs/train/trajgen/traj_amass.yaml task amass_traj
+```
+Train the motion model:
+```bash
+python tools/train.net -c configs/train/motiongen/motion_amass.yaml task amass_motion
+```
+The outputs and models will be saved in ```out/train/```
+## Finetune on the HUMANISE dataset
+Train the trajectory model:
+```bash
+python tools/train.net -c configs/train/trajgen/traj_humanise.yaml task humanise_traj resume True resume_model_dir out/train/amass_traj/model
+```
+Train the motion model:
+```bash
+python tools/train.net -c configs/train/motiongen/motion_humanise.yaml task humanise_motion resume True resume_model_dir out/train/amass_motion/model
+```
 # Citation
 
 ```

diff --git a/configs/dataset/amass.yaml b/configs/dataset/amass.yaml
@@ -0,0 +1,22 @@
+train_dat:
+  name: AMASS
+  limit_size: -1
+  split: train
+
+val_dat:
+  name: AMASS
+  limit_size: -1
+  split: test
+
+test_dat:
+  name: AMASS
+  limit_size: -1
+  split: test
+
+dat_cfg:
+  amass_root: data/amass_preprocess
+  max_motion_len: 120
+  num_scene_points: 1024
+  interval: 30
+  sample_data_interval: 1
+  preload: True
diff --git a/configs/train/default.yaml b/configs/train/default.yaml
@@ -0,0 +1,45 @@
+auto_config: []
+
+trainer:
+  name: default
+wrapper_cfg: 
+  name: default
+  pre_methods: []
+  post_methods: []
+  sup_methods: []
+  vis_methods: []
+
+loss_weights:
+  place_holder: 0
+
+metrics: []
+
+# ====== Train/Val/Test dataset/dataloader settings ====== #
+train:
+  epoch: 200
+  batch_size: &batch_size 256
+  shuffle: True
+  num_workers: 2
+  optimizer:
+    optim: adam
+    lr: 0.
+    canonical_lr: 1.e-4
+    canonical_bs: *batch_size
+    weight_decay: 0.0
+    adamw_weight_decay: 0.01
+  scheduler:
+    type: multi_step
+    milestones: [50, 100, 150]  # if epoch == 200
+    gamma: 0.5
+
+val:
+  epoch: 1
+  batch_size: 64
+  shuffle: False
+  num_workers: 2
+
+test:
+  epoch: 1
+  batch_size: 64
+  shuffle: False
+  num_workers: 2
diff --git a/configs/train/motiongen/motion_amass.yaml b/configs/train/motiongen/motion_amass.yaml
@@ -0,0 +1,57 @@
+task: auto
+
+dataset_cfg_path: "configs/dataset/amass.yaml"
+coord: &coord az
+
+wrapper_cfg:
+  name: MotionDiffuserWrapper
+  smplx_model_type: amass
+  pre_methods: []
+  normalizer:
+    name: NormalizerPoseMotion
+    file: out/release/normalize/amass_az_humanise_oc_transl_orient6d_pose6d.pkl
+
+eval_ep: 20
+metrics: ['recon_localpose', 'recon_trans', 'recon_orient']
+
+loss_weights:
+  recon_trans: 1.0
+  recon_orient_6d: 1.0
+  recon_pose_6d: 10.0
+
+train:
+  epoch: 200
+  batch_size: &batch_size 256
+  optimizer:
+    canonical_bs: *batch_size
+
+net_cfg:
+  coord: *coord
+  repr: motion # traj
+  name: DiffuserNetwork
+  k_sample: 10
+  diffuser: 
+    name: ObserConditionalDDPM
+    timesteps: 200
+    pred_type: pred_x0
+    obser: false
+    schedule_cfg:
+      beta: [0.0001, 0.01]
+      beta_schedule: cosine
+      s: 0.008
+  model:
+    name: MotionFromSceneTextTrajVoxelV0
+    d_l: 120
+    d_x: 135
+    d_betas: 10
+    env_sensor:
+      name: EnvSensor
+      voxel_dim: 8
+      radius: 2.0
+    target_sensor:
+      name: TargetSensor
+      voxel_dim: 8
+    traj_sensor:
+      name: TrajSensor
+      voxel_dim: 8
+      radius: 1.0
diff --git a/configs/train/motiongen/motion_humanise.yaml b/configs/train/motiongen/motion_humanise.yaml
@@ -0,0 +1,68 @@
+task: auto
+
+# resume: auto
+coord: &coord oc
+resume_model_dir: auto
+
+dataset_cfg_path: configs/dataset/humanise_motion.yaml
+
+wrapper_cfg:
+  name: MotionDiffuserWrapper
+  smplx_model_type: humanise
+  pre_methods: ['clip_text']
+  normalizer:
+    name: NormalizerPoseMotion
+    file: out/release/normalize/amass_az_humanise_oc_transl_orient6d_pose6d.pkl
+
+metrics: ['recon_localpose', 'recon_trans', 'recon_orient']
+
+loss_weights:
+  recon_trans: 1.0
+  recon_orient_6d: 1.0
+  recon_pose_6d: 10.0
+
+train:
+  epoch: 400
+  batch_size: 128
+  optimizer:
+    canonical_bs: 128
+  scheduler:
+    type: multi_step
+    milestones: [50, 100, 150, 250, 300, 350]  # if epoch == 200
+
+net_cfg:
+  coord: *coord
+  repr: motion
+  name: DiffuserNetwork
+  k_sample: 10
+  diffuser:
+    name: ObserConditionalDDPM
+    timesteps: 200
+    pred_type: pred_x0
+    obser: False
+    obser_type: start_motion
+    schedule_cfg:
+      beta: [0.0001, 0.01]
+      beta_schedule: cosine
+      s: 0.008
+  model:
+    name: MotionFromSceneTextTrajVoxelV0
+    d_l: 120
+    d_x: 135
+    d_betas: 10
+    env_sensor:
+      name: EnvSensor
+      voxel_dim: 8
+      radius: 2.0
+    target_sensor:
+      name: TargetSensor
+      voxel_dim: 8
+    traj_sensor:
+      name: TrajSensor
+      voxel_dim: 8
+      radius: 1.0
+
+  optimizer:
+    name: default
+  planner:
+    name: default
diff --git a/configs/train/trajgen/traj_amass.yaml b/configs/train/trajgen/traj_amass.yaml
@@ -0,0 +1,52 @@
+task: auto
+
+dataset_cfg_path: "configs/dataset/amass.yaml"
+
+wrapper_cfg:
+  name: MotionDiffuserWrapper
+  smplx_model_type: amass
+  pre_methods: []
+  normalizer:
+    name: NormalizerPoseMotion
+    file: out/release/normalize/amass_az_humanise_oc_transl_orient6d_pose6d.pkl
+
+eval_ep: 20
+metrics: ['recon_trans', 'recon_orient']
+
+loss_weights:
+  recon_trans: 1.0
+  recon_orient_6d: 1.0
+
+train:
+  epoch: 200
+  batch_size: 256
+  optimizer:
+    canonical_bs: 256
+
+net_cfg:
+  coord: az
+  repr: traj
+  name: DiffuserNetwork
+  k_sample: 10
+  diffuser:
+    name: ObserConditionalDDPM
+    timesteps: 200
+    pred_type: pred_x0
+    obser: false
+    obser_type: start_traj
+    schedule_cfg:
+      beta: [0.0001, 0.01]
+      beta_schedule: cosine
+      s: 0.008
+  model:
+    name: TrajFromSceneTextVoxelV0
+    d_l: 120
+    d_x: 9
+    d_betas: 10
+    env_sensor:
+      name: EnvSensor
+      voxel_dim: 8
+      radius: 2.0
+    target_sensor:
+      name: TargetSensor
+      voxel_dim: 8
diff --git a/configs/train/trajgen/traj_humanise.yaml b/configs/train/trajgen/traj_humanise.yaml
@@ -0,0 +1,56 @@
+task: auto
+
+# resume: auto
+resume_model_dir: auto
+dataset_cfg_path: configs/dataset/humanise_motion.yaml
+
+wrapper_cfg:
+  name: MotionDiffuserWrapper
+  smplx_model_type: humanise
+  pre_methods: ['clip_text']
+  normalizer:
+    name: NormalizerPoseMotion
+    file: out/release/normalize/amass_az_humanise_oc_transl_orient6d_pose6d.pkl
+
+metrics: ['recon_trans', 'recon_orient']
+
+loss_weights:
+  recon_trans: 1.0
+  recon_orient_6d: 1.0
+
+train:
+  epoch: 400
+  batch_size: 256
+  optimizer:
+    canonical_bs: 256
+  scheduler:
+    type: multi_step
+    milestones: [50, 100, 150, 250, 300, 350]
+
+net_cfg:
+  coord: oc
+  repr: traj
+  name: DiffuserNetwork
+  k_sample: 10
+  diffuser:
+    name: ObserConditionalDDPM
+    timesteps: 200
+    pred_type: pred_x0
+    obser: false
+    obser_type: start_traj
+    schedule_cfg:
+      beta: [0.0001, 0.01]
+      beta_schedule: cosine
+      s: 0.008
+  model:
+    name: TrajFromSceneTextVoxelV0
+    d_l: 120
+    d_x: 9
+    d_betas: 10
+    env_sensor:
+      name: EnvSensor
+      voxel_dim: 8
+      radius: 2.0
+    target_sensor:
+      name: TargetSensor
+      voxel_dim: 8
diff --git a/lib/config/config.py b/lib/config/config.py
@@ -26,9 +26,11 @@ def make_cfg(args):
     default_cfg_path = f"configs/{args.cfg_file.split('/')[1]}/default.yaml"
     if os.path.exists(default_cfg_path):
         cfg.merge_from_file(default_cfg_path)
+        cfg.merge_from_file(args.cfg_file)
         if 'dataset_cfg_path' in cfg.keys():
             cfg.merge_from_file(cfg.dataset_cfg_path)
-        cfg.merge_from_file(args.cfg_file)
+            cfg.merge_from_file(args.cfg_file)
+        cfg.merge_from_list(getattr(args, 'opts', []))
         # dirs
         if cfg.record_dir == 'auto':
             cfg.record_dir = f'out/train/{cfg.task}'
@@ -41,7 +43,7 @@ def make_cfg(args):
         logger.warning('overwrite gpus and resume!')
         cfg.gpus = [0]
         cfg.resume = True
-    cfg.merge_from_list(getattr(args, 'opts', []))
+        cfg.merge_from_list(getattr(args, 'opts', []))
     cfg.is_train = not args.is_test
 
     # 1. Auto config devices