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,96 @@
+import torch
+import wandb
+from tqdm import tqdm
+from trainer.train_builder import TRAINER
+
+
+@TRAINER.register("OAD")
+def train_one_epoch(trainloader, model, criterion, optimizer, scaler, epoch, use_wandb = False, writer=None, scheduler=None):
+    epoch_loss = 0
+    for it, (vid_name, rgb_input, flow_input, target) in enumerate(tqdm(trainloader, desc=f'Epoch:{epoch} Training', postfix=f'lr: {optimizer.param_groups[0]["lr"]:.7f}')):
+        rgb_input, flow_input, target = rgb_input.cuda(), flow_input.cuda(), target.cuda()
+        model.train()
+        if scaler != None:
+            with torch.cuda.amp.autocast():    
+                out_dict = model(rgb_input, flow_input) 
+                loss = criterion(out_dict, target)
+            optimizer.zero_grad(set_to_none=True)
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+            # scaler is None as argparse by default
+        else:
+            out_dict = model(rgb_input, flow_input)
+            loss = criterion(out_dict, target)
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+        
+        epoch_loss += loss.item()
+        if use_wandb == True:
+            wandb.log({"Train Loss": loss.item()})
+            
+        if writer != None:
+            writer.add_scalar("Train Loss", loss.item(), it+epoch*len(trainloader))
+    return epoch_loss
+
+@TRAINER.register("OADTR")
+def train_one_epoch_oadtr(trainloader, model, criterion, optimizer, scaler, epoch, use_wandb = False, writer=None, scheduler=None):
+    epoch_loss = 0
+    for it, (vid_name, rgb_input, flow_input, enc_target, distance_target, class_h_target, dec_target) in enumerate(tqdm(trainloader, desc=f'Epoch:{epoch} Training', postfix=f'lr: {optimizer.param_groups[0]["lr"]:.7f}')):
+        rgb_input, flow_input, enc_target, distance_target, class_h_target = rgb_input.cuda(), flow_input.cuda(), enc_target.cuda(), distance_target.cuda(), class_h_target.cuda()
+        model.train()
+        if scaler != None:
+            with torch.cuda.amp.autocast():    
+                out_dict = model(rgb_input, flow_input) 
+                loss = criterion(out_dict, enc_target, dec_target)   
+            optimizer.zero_grad(set_to_none=True)
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            out_dict = model(rgb_input, flow_input) 
+            loss = criterion(out_dict, enc_target, dec_target)
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+
+        epoch_loss += loss.item()
+        if use_wandb == True:
+            wandb.log({"Train Loss": loss.item()})
+            
+        if writer != None:
+            writer.add_scalar("Train Loss", loss.item(), it+epoch*len(trainloader))
+            
+    return epoch_loss
+
+
+@TRAINER.register("ANTICIPATION")
+def ant_train_one_epoch(trainloader, model, criterion, optimizer, scaler, epoch, use_wandb = False, writer=None, scheduler=None):
+    epoch_loss = 0
+    for it, (vid_name, rgb_input, flow_input, target, ant_target) in enumerate(tqdm(trainloader, desc=f'Epoch:{epoch} Training', postfix=f'lr: {optimizer.param_groups[0]["lr"]:.7f}')):
+        rgb_input, flow_input, target, ant_target = rgb_input.cuda(), flow_input.cuda(), target.cuda(), ant_target.cuda()
+        model.train()
+        if scaler != None:
+            with torch.cuda.amp.autocast():
+                out_dict = model(rgb_input, flow_input)
+                loss = criterion(out_dict, target, ant_target)
+            optimizer.zero_grad(set_to_none=True)
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            out_dict = model(rgb_input, flow_input) 
+            loss = criterion(out_dict, target, ant_target)
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+        epoch_loss += loss.item()
+        
+        if use_wandb == True:
+            wandb.log({"Train Loss": loss.item()})
+            
+        if writer != None:
+            writer.add_scalar("Train Loss", loss.item(), it+epoch*len(trainloader))
+            
+    return epoch_loss