""" python train.py \ --train_audio ./data/synth_dataset/train/audio \ --train_bs ./data/synth_dataset/train/bs \ --val_audio ./data/synth_dataset/val/audio \ --val_bs ./data/synth_dataset/val/bs \ --bs_dim 24 \ --batch_size 4 \ --epochs 5 \ --save_dir ./data/ckpt_synth24 """ import argparse, os, time, math, json from pathlib import Path import numpy as np import librosa from tqdm import tqdm import torch from torch.utils.data import Dataset, DataLoader from torch.utils.tensorboard import SummaryWriter from a2fmodel import A2BModel # --------------------------------------------------------------------- # ============ 1. 数据集 ============ class FaceDataset(Dataset): """ 读取配对的 (wav, bs.npy) - wav 采样率统一到 16k - bs.npy 形状 [T, N], N = --bs_dim """ def __init__(self, audio_dir, bs_dir): self.audio_paths = sorted(Path(audio_dir).glob("*.wav")) self.bs_dir = Path(bs_dir) assert len(self.audio_paths) > 0, f"no wav in {audio_dir}" def __len__(self): return len(self.audio_paths) def __getitem__(self, idx): wav_p = self.audio_paths[idx] bs_p = self.bs_dir / f"{wav_p.stem}.npy" assert bs_p.exists(), f"missing {bs_p}" audio, _ = librosa.load(wav_p, sr=16000) # (L,) audio = torch.from_numpy(audio).float().unsqueeze(0) # (1, L) bs = torch.from_numpy(np.load(bs_p)).float() # (T, N) return {"audio": audio, "bs": bs} # ============ 2. collate（变长 pad） ============ def collate_fn(batch): # --------- audio pad ---------- max_a = max(b["audio"].shape[1] for b in batch) audios = torch.zeros(len(batch), 1, max_a) for i, b in enumerate(batch): L = b["audio"].shape[1] audios[i, 0, :L] = b["audio"] # --------- bs pad ------------- max_t = max(b["bs"].shape[0] for b in batch) N = batch[0]["bs"].shape[1] bss = torch.zeros(len(batch), max_t, N) lens = torch.zeros(len(batch), dtype=torch.long) for i, b in enumerate(batch): T = b["bs"].shape[0] bss[i, :T] = b["bs"] lens[i] = T # A2BModel 期望的字段（3 份音频/target 可直接复用） data = { "input11": audios, # (B,1,L) "input12": audios.clone(), "input21": audios.clone(), "target11": bss, # (B,T,N) "target12": bss.clone(), "target21": bss.clone(), "level": torch.zeros(len(batch), dtype=torch.long), # 未使用，默认全 0 "person": torch.zeros(len(batch), dtype=torch.long), # 未使用，默认全 0 "lengths": lens, # 每条样本的有效 T } return data # ============ 3. 损失 ============ def mse_loss(pred, gt, mask): """ pred/gt: (B, T, N) mask: (B, T) (1=valid, 0=pad) """ diff = (pred - gt) ** 2 # (B,T,N) diff = diff.sum(-1) # (B,T) diff = diff * mask return diff.sum() / mask.sum().clamp_min(1.0) # ============ 4. 训练 / 验证 ============ def train_one_epoch(model, loader, optimizer, device): model.train() total_loss = 0.0 for batch in tqdm(loader, desc="Train", leave=False): for k in batch: if isinstance(batch[k], torch.Tensor): batch[k] = batch[k].to(device, non_blocking=True) bs_pred1, bs_pred2, _ = model(batch) # (B,T,N) tgt = batch["target11"] # (B,T,N) # 构造 mask (B,T) B, T, _ = tgt.shape mask = (torch.arange(T, device=device)[None, :] < batch["lengths"][:, None]).float() loss = mse_loss(bs_pred1, tgt, mask) + mse_loss(bs_pred2, tgt, mask) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() return total_loss / max(len(loader), 1) @torch.no_grad() def evaluate(model, loader, device): model.eval() total_loss = 0.0 for batch in loader: for k in batch: if isinstance(batch[k], torch.Tensor): batch[k] = batch[k].to(device, non_blocking=True) bs_pred1, bs_pred2, _ = model(batch) tgt = batch["target11"] B, T, _ = tgt.shape mask = (torch.arange(T, device=device)[None, :] < batch["lengths"][:, None]).float() loss = mse_loss(bs_pred1, tgt, mask) + mse_loss(bs_pred2, tgt, mask) total_loss += loss.item() return total_loss / max(len(loader), 1) # ============ 5. CLI / 入口 ============ def main(): parser = argparse.ArgumentParser() # ---------- 数据路径 ---------- parser.add_argument("--train_audio", required=True, type=str) parser.add_argument("--train_bs", required=True, type=str) parser.add_argument("--val_audio", type=str, default=None) parser.add_argument("--val_bs", type=str, default=None) # ---------- 训练超参 ---------- parser.add_argument("--epochs", type=int, default=20) parser.add_argument("--batch_size", type=int, default=2) parser.add_argument("--lr", type=float, default=1e-4) parser.add_argument("--device", type=str, default="cuda:0") parser.add_argument("--save_dir", type=str, default="./checkpoints") # ---------- A2BModel 超参 ---------- parser.add_argument("--bs_dim", type=int, default=52) # 可改为你的 DOF 数 parser.add_argument("--feature_dim", type=int, default=832) parser.add_argument("--period", type=int, default=30) parser.add_argument("--max_seq_len", type=int, default=5000) parser.add_argument("--emo_guide", action="store_true") args = parser.parse_args() # ---------- 目录 / 日志 ---------- Path(args.save_dir).mkdir(parents=True, exist_ok=True) with open(Path(args.save_dir) / "hparams.json", "w") as f: json.dump(vars(args), f, indent=2, ensure_ascii=False) writer = SummaryWriter(log_dir=os.path.join(args.save_dir, "tb")) device = torch.device(args.device if torch.cuda.is_available() else "cpu") # ---------- DataLoader ---------- train_set = FaceDataset(args.train_audio, args.train_bs) train_loader = DataLoader( train_set, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True, collate_fn=collate_fn ) val_loader = None if args.val_audio and args.val_bs: val_set = FaceDataset(args.val_audio, args.val_bs) val_loader = DataLoader( val_set, batch_size=args.batch_size, shuffle=False, num_workers=2, pin_memory=True, collate_fn=collate_fn ) # ---------- 模型 ---------- model_cfg = argparse.Namespace( bs_dim=args.bs_dim, feature_dim=args.feature_dim, period=args.period, max_seq_len=args.max_seq_len, emo_guide=args.emo_guide, device=args.device, # A2BModel 内部会用 .to(self.device) batch_size=args.batch_size, ) model = A2BModel(model_cfg).to(device) # ---------- 优化器 / 学习率计划 ---------- optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs) # ---------- 训练循环 ---------- best_val = math.inf for epoch in range(1, args.epochs + 1): t0 = time.time() train_loss = train_one_epoch(model, train_loader, optimizer, device) writer.add_scalar("loss/train", train_loss, epoch) msg = f"[{epoch:03d}/{args.epochs}] train {train_loss:.4f}" if val_loader is not None: val_loss = evaluate(model, val_loader, device) writer.add_scalar("loss/val", val_loss, epoch) msg += f" | val {val_loss:.4f}" if val_loss < best_val: best_val = val_loss torch.save(model.state_dict(), Path(args.save_dir) / "best.pth") # 保存最近模型 torch.save(model.state_dict(), Path(args.save_dir) / "last.pth") scheduler.step() dt = time.time() - t0 print(msg + f" | time {dt:.1f}s") writer.close() if best_val < math.inf: print(f"Training finished. Best val loss = {best_val:.4f}") else: print("Training finished.") if __name__ == "__main__": main()