File size: 5,400 Bytes
a89a522 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | """Root model training script using synthetic data.
Demonstrates how the root backbone training pipeline works without
requiring the actual motion dataset. Loads the saved model config from
the checkpoint directory and trains on randomly generated motion tensors.
The root model does not require a pretrained VQVAE — it directly
predicts continuous root motion values.
Usage:
python scripts/train_root.py --max_steps 100
"""
import argparse
import copy
import os
import pytorch_lightning as pl
import torch
from hydra.utils import instantiate
from omegaconf import DictConfig, OmegaConf, open_dict
from torch.utils.data import DataLoader
from motionbricks.data.synthetic_dataset import SyntheticMotionDataset, collate_batch
from motionbricks.helper.pl_util import load_motion_rep
def load_config(result_dir: str, max_steps: int):
"""Load and patch hparams.yaml for single-GPU training."""
version_dir = os.path.join(result_dir, "motionbricks_root", "version_1")
hparams_path = os.path.join(version_dir, "hparams.yaml")
conf = OmegaConf.load(hparams_path)
with open_dict(conf):
# resolve data paths to the version directory (where skeleton/stats live)
conf.data = {"folder": version_dir, "text_embeddings": None}
conf.skeleton.folder = os.path.join(version_dir, "skeleton")
conf.motion_rep.stats.folder = os.path.join(version_dir, "stats", "motion")
# single-GPU training overrides
conf.trainer.devices = 1
conf.trainer.num_nodes = 1
conf.trainer.max_steps = max_steps
conf.trainer.accelerator = "auto"
conf.trainer.strategy = "auto"
conf.trainer.enable_progress_bar = True
conf.trainer.log_every_n_steps = 10
conf.trainer.val_check_interval = max_steps
conf.trainer.num_sanity_val_steps = 0
# resolve ${trainer.max_steps} in scheduler
conf.model.scheduler.num_training_steps = max_steps
# remove keys with unresolvable ${hydra:...} interpolations
conf.id = "synthetic"
conf.run_dir = "."
conf.out_dir = result_dir
# resolve all ${} interpolations, then re-wrap as DictConfig
resolved = OmegaConf.to_container(conf, resolve=True)
conf = OmegaConf.create(resolved)
return conf, version_dir
def main():
parser = argparse.ArgumentParser(description="Root model training")
parser.add_argument("--result_dir", type=str, default="./out",
help="Directory containing pretrained checkpoints")
parser.add_argument("--max_steps", type=int, default=200,
help="Number of training steps")
parser.add_argument("--batch_size", type=int, default=8,
help="Batch size")
parser.add_argument("--num_samples", type=int, default=500,
help="Number of synthetic samples in dataset")
parser.add_argument("--seed", type=int, default=42)
args = parser.parse_args()
pl.seed_everything(args.seed)
conf, version_dir = load_config(args.result_dir, args.max_steps)
# instantiate skeleton and motion representation
motion_rep = load_motion_rep(conf)
feat_dim = len(motion_rep.indices['all'])
# create synthetic dataset
dataset = SyntheticMotionDataset(
feat_dim=feat_dim,
num_samples=args.num_samples,
min_frames=200,
max_frames=400,
)
dataloader = DataLoader(
dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=2,
collate_fn=collate_batch,
persistent_workers=True,
)
# instantiate networks and model
model_conf = copy.deepcopy(conf.model)
with open_dict(model_conf):
# instantiate backbone network (needs full motion_rep for dual_rep access)
backbone_net = instantiate(
model_conf.backbone_network,
motion_rep=motion_rep,
_recursive_=False,
)
# build optimizer and scheduler as partials
optimizer_fn = instantiate(model_conf.optimizer)
scheduler_fn = instantiate(model_conf.scheduler) if model_conf.scheduler else None
model = instantiate(
model_conf,
pose_vqvae_network=None,
root_vqvae_network=None,
backbone_network=backbone_net,
motion_rep=motion_rep,
optimizer=optimizer_fn,
scheduler=scheduler_fn,
_recursive_=False,
)
# create trainer (no callbacks)
trainer = pl.Trainer(
max_steps=conf.trainer.max_steps,
devices=conf.trainer.devices,
num_nodes=conf.trainer.num_nodes,
accelerator=conf.trainer.accelerator,
strategy=conf.trainer.strategy,
precision=conf.trainer.precision,
gradient_clip_val=conf.trainer.gradient_clip_val,
enable_progress_bar=conf.trainer.enable_progress_bar,
log_every_n_steps=conf.trainer.log_every_n_steps,
num_sanity_val_steps=0,
enable_checkpointing=False,
logger=False,
)
print(f"Starting root model training for {args.max_steps} steps...")
print(f" Feature dim: {feat_dim}")
print(f" Batch size: {args.batch_size}")
print(f" Dataset size: {args.num_samples}")
trainer.fit(model, train_dataloaders=dataloader)
print("Training complete.")
if __name__ == "__main__":
main()
|