Refactored code for AddBiomechanics and previous datasets

Author: Starfractor

Dockerfile

@@ -1,11 +1,11 @@
 # Use an official Ubuntu runtime as a parent image
-FROM ubuntu:22.04
+FROM ubuntu:25.04
 
 # Set the working directory
 WORKDIR /P-BIGE
 
 # Install necessary dependencies
-RUN apt-get update && apt-get install -y wget git htop xvfb
+RUN apt-get update && apt-get install -y wget git htop xvfb build-essential 
 
 # Install Miniconda
 RUN MINICONDA_INSTALLER_SCRIPT=Miniconda3-py38_23.1.0-1-Linux-x86_64.sh && \
@@ -27,10 +27,6 @@ RUN conda env create -f environment.yml
 # Activate the conda environment for subsequent RUN commands
 SHELL ["conda", "run", "-n", "P-BIGE", "/bin/bash", "-c"]
 
-# Download the model and extractor
-RUN bash dataset/prepare/download_model.sh && \
-    bash dataset/prepare/download_extractor.sh
-
 # Install additional Python packages (if needed)
 RUN pip install --user ipykernel polyscope easydict trimesh
 RUN pip install --user --force-reinstall numpy==1.22.0

dataset/dataset_MOT_segmented.py

@@ -205,7 +205,7 @@ def __getitem__(self, item):
             return subsequences, subsequence_lengths, names
 
 class AddBiomechanicsDataset(data.Dataset):
-    def __init__(self, window_size=64, unit_length=4, mode='train', data_dir='/home/mnt/data/addb_dataset_publication'):
+    def __init__(self, window_size=64, unit_length=4, mode='train', data_dir='/home/kingn450/Datasets/addb_dataset_publication'):
         self.window_size = window_size
         self.unit_length = unit_length
         self.data_dir = data_dir

environment.yml

@@ -10,7 +10,6 @@ dependencies:
   - bzip2=1.0.8=h7b6447c_0
   - ca-certificates=2021.7.5=h06a4308_1
   - certifi=2021.5.30=py38h06a4308_0
-  - cudatoolkit=10.1.243=h6bb024c_0
   - ffmpeg=4.3=hf484d3e_0
   - freetype=2.10.4=h5ab3b9f_0
   - gmp=6.2.1=h2531618_2
@@ -56,8 +55,8 @@ dependencies:
   - six=1.16.0=pyhd3eb1b0_0
   - sqlite=3.36.0=hc218d9a_0
   - tk=8.6.10=hbc83047_0
-  - torchaudio=0.8.1=py38
-  - torchvision=0.9.1=py38_cu101
+  - torchaudio=2.3.0
+  - torchvision
   - typing_extensions=3.10.0.0=pyh06a4308_0
   - wheel=0.37.0=pyhd3eb1b0_0
   - xz=5.2.5=h7b6447c_0

output_test/TEST_RUN/events.out.tfevents.1746785869.DESKTOP-42KLHKG.66316.0

@@ -1,41 +1,32 @@
-import os
+import deepspeed
 import json
-# from osim_sequence import OSIMSequence,load_osim
+import models.vqvae as vqvae
+import nimblephysics as nimble
+import options.option_vq as option_vq
+import os
 import torch
+import torch.distributed as dist
 import torch.optim as optim
-from torch.utils.tensorboard import SummaryWriter
-from tqdm import tqdm
-
-import models.vqvae as vqvae
+import utils.eval_trans as eval_trans
 import utils.losses as losses 
-import options.option_vq as option_vq
 import utils.utils_model as utils_model
+import warnings
 from dataset import dataset_MOT_MCS, dataset_TM_eval, dataset_MOT_segmented
-import utils.eval_trans as eval_trans
-from options.get_eval_option import get_opt
 from models.evaluator_wrapper import EvaluatorModelWrapper
-import warnings
-warnings.filterwarnings('ignore')
+from options.get_eval_option import get_opt
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
 from utils.word_vectorizer import WordVectorizer
-import nimblephysics as nimble
-import deepspeed
-
+warnings.filterwarnings('ignore')
 
 def update_lr_warm_up(optimizer, nb_iter, warm_up_iter, lr):
-
     current_lr = lr * (nb_iter + 1) / (warm_up_iter + 1)
     for param_group in optimizer.param_groups:
         param_group["lr"] = current_lr
-
     return optimizer, current_lr
 
 def get_foot_losses(motion, y_translation=0.0,feet_threshold=0.01):
-    # y_translation = 0.0
     min_height, idx = motion[..., 1].min(dim=-1)
-
-    # y_translation = -min_height.median() # Change reference to median  (Other set of experiments determine this. See paper)
-
-    # print(min_height,idx,motion[..., 1].shape)
     min_height = min_height + y_translation
     pn = -torch.minimum(min_height, torch.zeros_like(min_height))  # penetration
     pn[pn < feet_threshold] = 0.0
@@ -64,16 +55,25 @@ def get_foot_losses(motion, y_translation=0.0,feet_threshold=0.01):
 
     return loss_pn.sum()/bs, loss_fl.sum()/bs, loss_sk.sum()/bs
 
-##### ---- Exp dirs ---- #####
+# --- Robust device and distributed/deepspeed setup ---
 args = option_vq.get_args_parser()
 torch.manual_seed(args.seed)
+
+args.local_rank = int(os.environ.get("LOCAL_RANK", 0))
+world_size = int(os.environ.get("WORLD_SIZE", 1))
+
 if torch.cuda.is_available():
     torch.cuda.set_device(args.local_rank)
+    if dist.is_available() and not dist.is_initialized() and world_size > 1:
+        dist.init_process_group(backend="nccl")
+    device = torch.device(f"cuda:{args.local_rank}")
+else:
+    device = torch.device("cpu")
 
 args.out_dir = os.path.join(args.out_dir, f'{args.exp_name}')
 os.makedirs(args.out_dir, exist_ok = True)
 
-##### ---- Logger ---- #####
+# Logger
 logger = utils_model.get_logger(args.out_dir)
 writer = SummaryWriter(args.out_dir)
 logger.info(json.dumps(vars(args), indent=4, sort_keys=True))
@@ -83,54 +83,37 @@ def get_foot_losses(motion, y_translation=0.0,feet_threshold=0.01):
 if args.dataname == 'kit' : 
     dataset_opt_path = 'checkpoints/kit/Comp_v6_KLD005/opt.txt'  
     args.nb_joints = 21
-    
+if args.dataname == 'addb' :
+    args.nb_joints = 23
 else :
     dataset_opt_path = 'checkpoints/t2m/Comp_v6_KLD005/opt.txt'
     args.nb_joints = 22
 
-args.nb_joints = 23 # fixed issues
-
 logger.info(f'Training on {args.dataname}, motions are with {args.nb_joints} joints')
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-wrapper_opt = get_opt(dataset_opt_path, device)
-#eval_wrapper = EvaluatorModelWrapper(wrapper_opt)
-
-
-##### ---- Dataloader ---- #####
-# train_loader = dataset_MOT_MCS.DATALoader(args.dataname,
-#                                         args.batch_size,
-#                                         window_size=args.window_size,
-#                                         unit_length=2**args.down_t)
-
 train_loader = dataset_MOT_segmented.addb_data_loader(
     window_size=args.window_size,
     unit_length=2**args.down_t,
     batch_size=args.batch_size,
     mode=args.dataname
 )
 
-# train_loader_iter = dataset_MOT_MCS.cycle(train_loader)
 train_loader_iter = dataset_MOT_segmented.cycle(train_loader)
 
-# val_loader = dataset_TM_eval.DATALoader(args.dataname, False,
-#                                         32,
-#                                         w_vectorizer,
-#                                         unit_length=2**args.down_t)
-
-##### ---- Network ---- #####
-net = vqvae.HumanVQVAE(args, ## use args to define different parameters in different quantizers
-                       args.nb_code,
-                       args.code_dim,
-                       args.output_emb_width,
-                       args.down_t,
-                       args.stride_t,
-                       args.width,
-                       args.depth,
-                       args.dilation_growth_rate,
-                       args.vq_act,
-                       args.vq_norm)
-
+# Setup VQ-VAE model
+net = vqvae.HumanVQVAE(
+    args,
+    args.nb_code,
+    args.code_dim,
+    args.output_emb_width,
+    args.down_t,
+    args.stride_t,
+    args.width,
+    args.depth,
+    args.dilation_growth_rate,
+    args.vq_act,
+    args.vq_norm
+)
 
 if args.resume_pth : 
     logger.info('loading checkpoint from {}'.format(args.resume_pth))
@@ -139,67 +122,47 @@ def get_foot_losses(motion, y_translation=0.0,feet_threshold=0.01):
 net.train()
 net.to(device)
 
-##### ---- Optimizer & Scheduler ---- #####
+# Optimizer and scheduler setup
 optimizer = optim.AdamW(net.parameters(), lr=args.lr, betas=(0.9, 0.99), weight_decay=args.weight_decay)
 scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.lr_scheduler, gamma=args.gamma)
 
-deepspeed_config = {
-    "train_micro_batch_size_per_gpu": args.batch_size,
-    "optimizer": {
-        "type": "AdamW",
-        "params": {
-            "lr": args.lr,
-            "betas": [
-                0.9,
-                0.99
-            ],
-            "weight_decay":args.weight_decay
-        }
-    },
-    "gradient_accumulation_steps": 1,
-    # "fp16": {
-    #     "enabled": True
-    # },
-    "zero_optimization": {
-        "stage": 0
+# Run deepspeed if avaliable (2+ GPUs)
+if torch.cuda.is_available() and world_size > 1:
+    deepspeed_config = {
+        "train_micro_batch_size_per_gpu": args.batch_size,
+        "optimizer": {
+            "type": "AdamW",
+            "params": {
+                "lr": args.lr,
+                "betas": [0.9, 0.99],
+                "weight_decay": args.weight_decay
+            }
+        },
+        "gradient_accumulation_steps": 1,
+        "zero_optimization": {"stage": 0}
     }
-}
-net, optimizer, _, _ = deepspeed.initialize(model=net, optimizer=optimizer, args=args, config_params=deepspeed_config)
+    net, optimizer, _, _ = deepspeed.initialize(
+        model=net,
+        optimizer=optimizer,
+        args=args,
+        config_params=deepspeed_config
+    )
+else:
+    logger.info("Running without DeepSpeed (single GPU or CPU).")
 
 Loss = losses.ReConsLoss(args.recons_loss, args.nb_joints)
 
-##### ------ warm-up ------- #####
+# Warm up
 avg_recons, avg_perplexity, avg_commit, avg_temporal = 0., 0., 0., 0.
-
 for nb_iter in range(1, args.warm_up_iter):
-    
     optimizer, current_lr = update_lr_warm_up(optimizer, nb_iter, args.warm_up_iter, args.lr)
-    
     gt_motion,_, names = next(train_loader_iter)
-    gt_motion = gt_motion.to(device).float() # (bs, 64, dim)
+    gt_motion = gt_motion.to(device).float() 
 
     pred_motion, loss_commit, perplexity = net(gt_motion)
     loss_motion = Loss(pred_motion, gt_motion)
-    
-    loss_temp = torch.mean((pred_motion[:,1:,:]-pred_motion[:,:-1,:])**2)
-    
-    # loss_vel = Loss.forward_vel(pred_motion, gt_motion)
-    # loss_pn, loss_fl, loss_sk = get_foot_losses(pred_motion)
-    # print(loss_pn, loss_fl, loss_sk)
-    
-    # # hip flexion 7, 15
-    # hip_flexion_l = -pred_motion[:,:,7].max(dim=1).values.mean()
-    # hip_flexion_r = -pred_motion[:,:,15].max(dim=1).values.mean()
-    # # knee angle 10, 18
-    # knee_angle_l = -pred_motion[:,:,10].max(dim=1).values.mean()
-    # knee_angle_r = -pred_motion[:,:,18].max(dim=1).values.mean()
-    # # ankle angle 12, 20
-    # ankle_angle_l = -pred_motion[:,:,12].max(dim=1).values.mean()
-    # ankle_angle_r = -pred_motion[:,:,20].max(dim=1).values.mean()
-    
-    # print(hip_flexion_l, hip_flexion_r, knee_angle_l, knee_angle_r, ankle_angle_l, ankle_angle_r)
-    
-    loss = loss_motion + args.commit * loss_commit + 0.5 * loss_temp #+ args.loss_vel * loss_vel 
+    loss_temp = torch.mean((pred_motion[:,1:,:] - pred_motion[:,:-1,:])**2)
+    loss = loss_motion + args.commit * loss_commit + 0.5 * loss_temp 
 
     optimizer.zero_grad()
     loss.backward()
@@ -220,24 +183,17 @@ def get_foot_losses(motion, y_translation=0.0,feet_threshold=0.01):
 
         avg_recons, avg_perplexity, avg_commit, avg_temporal = 0., 0., 0., 0.
 
-##### ---- Training ---- #####
+# Training Loop
 avg_recons, avg_perplexity, avg_commit, avg_temporal = 0., 0., 0., 0.
 torch.save({'net' : net.state_dict()}, os.path.join(args.out_dir, 'warmup.pth'))
-# best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, writer, logger = eval_trans.evaluation_vqvae(args.out_dir, val_loader, net, logger, writer, 0, best_fid=1000, best_iter=0, best_div=100, best_top1=0, best_top2=0, best_top3=0, best_matching=100, eval_wrapper=eval_wrapper)
-
 for nb_iter in range(1, args.total_iter + 1):
-    
     gt_motion,_,_ = next(train_loader_iter)
-    gt_motion = gt_motion.to(device).float() # bs, nb_joints, joints_dim, seq_len
+    gt_motion = gt_motion.to(device).float() 
 
     pred_motion, loss_commit, perplexity = net(gt_motion)
     loss_motion = Loss(pred_motion, gt_motion)
     loss_temp = torch.mean((pred_motion[:,1:,:]-pred_motion[:,:-1,:])**2)
-    # loss_vel = Loss.forward_vel(pred_motion, gt_motion)
-    # loss_pn, loss_fl, loss_sk = get_foot_losses(pred_motion)
-    # print(loss_pn, loss_fl, loss_sk)
-    
-    loss = loss_motion + args.commit * loss_commit + 0.5 * loss_temp #+ args.loss_vel * loss_vel # Need to remove/change loss_vel since its not SMPL
+    loss = loss_motion + args.commit * loss_commit + 0.5 * loss_temp 
 
     optimizer.zero_grad()
     loss.backward()
@@ -261,16 +217,8 @@ def get_foot_losses(motion, y_translation=0.0,feet_threshold=0.01):
 
         logger.info(f"Train. Iter {nb_iter} : \t Commit. {avg_commit:.5f} \t PPL. {avg_perplexity:.2f} \t Recons.  {avg_recons:.5f} \t Temporal. {avg_temporal:.5f}")
 
-        avg_recons, avg_perplexity, avg_commit = 0., 0., 0.,
-    
-    if nb_iter % (10*args.eval_iter) == 0:
+        avg_recons, avg_perplexity, avg_commit, avg_temporal = 0., 0., 0., 0.
+
+    if nb_iter % (10 * args.eval_iter) == 0:
         torch.save({'net' : net.state_dict()}, os.path.join(args.out_dir, str(nb_iter) + '.pth'))
 
-    # if nb_iter % args.eval_iter==0 :
-    #     # The line `best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching,
-    #     # writer, logger = eval_trans.evaluation_vqvae(args.out_dir, val_loader, net, logger, writer,
-    #     # nb_iter, best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching,
-    #     # eval_wrapper=eval_wrapper)` is calling a function named `evaluation_vqvae` from the
-    #     # `eval_trans` module.
-    #     best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, writer, logger = eval_trans.evaluation_vqvae(args.out_dir, val_loader, net, logger, writer, nb_iter, best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, eval_wrapper=eval_wrapper)
-