Add causal person fusion and MuVAP stage-3 baselines

2026-09-10 06:54:22 +08:00
parent 0534c9d3d0
commit d426b4740a
225 changed files with 27716 additions and 119 deletions
@@ -0,0 +1,97 @@
+import os
+
+import wandb
+
+
+def parse_config(cfg):
+    modality = get_modality(cfg)
+    
+    cfg_filename = os.path.basename(cfg.OUTPUT_DIR)
+    
+    os.environ["WANDB__SERVICE_WAIT"] = "3000"
+    cfg_filename_abbr = cfg_filename.split('_1x_')[-1]
+    wandb_postfix = generate_wandb_postfix(cfg, modality)
+    scheduler_name = cfg.SOLVER.SCHEDULER.SCHEDULER_NAME
+    learning_rate = cfg.SOLVER.BASE_LR
+    weight_decay = cfg.SOLVER.WEIGHT_DECAY
+    
+    cfg_filename = cfg_filename.replace('long_512', f'long_{cfg.MODEL.LSTR.LONG_MEMORY_SECONDS}')
+    cfg_filename = cfg_filename.replace('work_8', f'work_{cfg.MODEL.LSTR.WORK_MEMORY_SECONDS}')
+    
+    cfg.OUTPUT_DIR = os.path.join(os.path.dirname(cfg.OUTPUT_DIR), cfg_filename)
+    
+    if cfg.SAVE_HPARAM:
+        # save different dir containing hparam
+        hparam_info = '_'.join([f"lr_{learning_rate}", f"wd_{weight_decay}", f"schd_{scheduler_name}"])
+        cfg.OUTPUT_DIR = os.path.join(cfg.OUTPUT_DIR, hparam_info, modality)
+        print(f"HPARAM: {hparam_info}")
+        print(f"cfg.OUTPUT_DIR: {cfg.OUTPUT_DIR}")
+            
+    wandb_name = cfg_filename_abbr + '_' + wandb_postfix
+    
+    if cfg.SAVE_HPARAM:
+        wandb_name += '_' + '_'.join([f"lr_{learning_rate}", f"wd_{weight_decay}", f"schd_{scheduler_name}"])
+
+    if cfg.USE_WANDB:
+        print("*** init wandb ***")
+        print(f"wandb_name: {wandb_name}")
+
+def get_modality(cfg):
+    modality = ''
+    if 'motion' in cfg.INPUT.MODALITY:
+        modality += 'A'
+    if 'visual' in cfg.INPUT.MODALITY:
+        modality += 'V'
+    if 'object' in cfg.INPUT.MODALITY:
+        modality += 'T'
+    elif cfg['INPUT']['MODALITY'] == 'twostream':
+        modality = 'AV'
+    
+    if modality == '':
+        raise ValueError(f"Unknown modality: {cfg['INPUT']['MODALITY']}")    
+    
+    return modality
+
+
+def generate_wandb_postfix(cfg, modality):
+    visual_feature_type = '_'.join(cfg['INPUT']['VISUAL_FEATURE'].split('_')[1:])
+    audio_feature_type = '_'.join(cfg['INPUT']['MOTION_FEATURE'].split('_')[1:])
+    text_feature_type = '_'.join(cfg['INPUT']['OBJECT_FEATURE'].split('_')[1:])
+    
+    if cfg.MODEL.LSTR.LONG_MEMORY_USE_PE == True:
+        model_name = 'LSTR_TESTRA'
+    else:
+        model_name = 'TESTRA'
+
+    if '_aa' in cfg.OUTPUT_DIR:
+        model_name += '_aa'
+    
+    name_parts = [
+        model_name,
+        cfg['DATA']['DATA_NAME'],
+        modality,        
+        # visual_feature_type,
+        # audio_feature_type,
+        # text_feature_type,
+        # f"epoch_{cfg['SOLVER']['NUM_EPOCHS']}"
+    ]
+    if cfg['MODEL']['CHECKPOINT']:
+        name_parts.append(f"ckpt_{os.path.basename(cfg['MODEL']['CHECKPOINT']).split('.')[0]}")
+    return '_'.join(name_parts)
+  
+def get_wandb_name(cfg, cfg_filename, modality):
+    cfg_filename_abbr = cfg_filename.split('_1x_')[-1]
+    wandb_postfix = generate_wandb_postfix(cfg, modality)
+    scheduler_name = cfg.SOLVER.SCHEDULER.SCHEDULER_NAME
+    learning_rate = cfg.SOLVER.BASE_LR
+    weight_decay = cfg.SOLVER.WEIGHT_DECAY
+    wandb_name = cfg_filename_abbr + '_' + wandb_postfix + '_' + '_'.join([f"lr_{learning_rate}", f"wd_{weight_decay}", f"schd_{scheduler_name}"])
+    return wandb_name
+
+def validate_cfg(cfg):
+    if cfg.DATA.NUM_CLASSES != len(cfg.DATA.CLASS_NAMES):
+        print(f"len(cfg.DATA.CLASS_NAMES): {len(cfg.DATA.CLASS_NAMES)} is not equal to cfg.DATA.NUM_CLASSES: {cfg.DATA.NUM_CLASSES}")
+        cfg.DATA.CLASS_NAMES = cfg.DATA.CLASS_NAMES[:cfg.DATA.NUM_CLASSES]
+        print(f"Adjusted cfg.DATA.CLASS_NAMES: {cfg.DATA.CLASS_NAMES}")
+        
+