chore: initial commit — unified project repo
Merged code repo (CompanionGuard-RL) into single project-level git. Reorganized root: docs/, reference/, experiments/, tmp/active|archives/. Gitignored: data/, checkpoints/, .venv, experiment logs, tmp/archives. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
630
旧方向信息/scripts/upload_train_script.py
Normal file
630
旧方向信息/scripts/upload_train_script.py
Normal file
@@ -0,0 +1,630 @@
|
||||
"""Upload train_d1.py and config files to server."""
|
||||
import paramiko, warnings
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
client = paramiko.SSHClient()
|
||||
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||||
client.connect('10.82.3.180', port=20083, username='root', password='m2dGcwyrhI', timeout=30)
|
||||
sftp = client.open_sftp()
|
||||
|
||||
ZSY = '/root/siton-data-2849d4ce327c4ccfb233ce33868fe7fe/zsy'
|
||||
PROJ = ZSY + '/multimodal_affect'
|
||||
|
||||
# ─── scripts/train/train_d1.py ────────────────────────────────────────────
|
||||
TRAIN_D1 = '''\
|
||||
#!/usr/bin/env python3
|
||||
# Phase 1 Direction 1 Training Script
|
||||
# Stage A: Supervised pretraining with noise-aware confidence estimation
|
||||
# Stage B: PPO-based adaptive fusion weight learning
|
||||
#
|
||||
# Launch:
|
||||
# torchrun --nproc_per_node=4 scripts/train/train_d1.py \\
|
||||
# --stage supervised --dataset IEMOCAP \\
|
||||
# --config configs/d1/stage_a.yaml \\
|
||||
# --output outputs/checkpoints/d1_stageA
|
||||
|
||||
import os, sys, argparse, yaml, time, logging
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
from torch.cuda.amp import GradScaler, autocast
|
||||
from sklearn.metrics import f1_score, accuracy_score
|
||||
import wandb
|
||||
|
||||
ZSY = os.environ.get("ZSY", "/root/siton-data-2849d4ce327c4ccfb233ce33868fe7fe/zsy")
|
||||
PROJ = os.path.join(ZSY, "multimodal_affect")
|
||||
sys.path.insert(0, PROJ)
|
||||
|
||||
from src.data.dataset import MultimodalDataset, get_dataloader
|
||||
from src.models.encoders import MultimodalEncoder
|
||||
from src.models.classifier import EmotionClassifier
|
||||
from src.rl.fusion_agent import ModalFusionAgent
|
||||
from src.rl.reward import compute_reward
|
||||
|
||||
|
||||
# ── Distributed helpers ───────────────────────────────────────────────────
|
||||
|
||||
def setup_ddp():
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", 0))
|
||||
dist.init_process_group("nccl")
|
||||
torch.cuda.set_device(local_rank)
|
||||
return local_rank, dist.get_rank(), dist.get_world_size()
|
||||
|
||||
def cleanup():
|
||||
dist.destroy_process_group()
|
||||
|
||||
def is_main(rank):
|
||||
return rank == 0
|
||||
|
||||
def all_reduce_mean(val, device):
|
||||
t = torch.tensor(float(val), device=device)
|
||||
dist.all_reduce(t, op=dist.ReduceOp.SUM)
|
||||
return (t / dist.get_world_size()).item()
|
||||
|
||||
def save_ckpt(state, path):
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
torch.save(state, path)
|
||||
|
||||
|
||||
def _noisy_batch(dataset, variant, indices, device):
|
||||
text = variant.get("text", dataset.text)
|
||||
audio = variant.get("audio", dataset.audio)
|
||||
vision = variant.get("vision", dataset.vision)
|
||||
return (
|
||||
torch.from_numpy(text[indices]).to(device),
|
||||
torch.from_numpy(audio[indices]).to(device),
|
||||
torch.from_numpy(vision[indices]).to(device),
|
||||
torch.from_numpy(dataset.labels[indices]).to(device),
|
||||
)
|
||||
|
||||
|
||||
def _confidence_targets(variant_name, batch_size, device):
|
||||
target = torch.full((batch_size, 3), 0.9, device=device)
|
||||
noisy_map = {
|
||||
"gaussian_light": (0, 1, 2),
|
||||
"gaussian_heavy": (0, 1, 2),
|
||||
"missing_audio": (1,),
|
||||
"missing_visual": (2,),
|
||||
"text_word_drop_30": (0,),
|
||||
"audio_masking_50": (1,),
|
||||
"realistic_mixed": (0, 1, 2),
|
||||
"audio_time_mask": (1,),
|
||||
}
|
||||
for idx in noisy_map.get(str(variant_name), (0, 1, 2)):
|
||||
target[:, idx] = 0.1
|
||||
return target
|
||||
|
||||
|
||||
# ── Evaluation ────────────────────────────────────────────────────────────
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate(encoder, classifier, loader, device, agent=None):
|
||||
encoder.eval()
|
||||
classifier.eval()
|
||||
if agent is not None:
|
||||
agent.eval()
|
||||
all_preds, all_labels = [], []
|
||||
for batch in loader:
|
||||
text = batch["text"].to(device)
|
||||
audio = batch["audio"].to(device)
|
||||
vision = batch["vision"].to(device)
|
||||
labels = batch["labels"].to(device)
|
||||
tf, af, vf, confs = encoder(text, audio, vision)
|
||||
if agent is not None:
|
||||
noise_est = audio.std(dim=-1, keepdim=True).sigmoid()
|
||||
state = torch.cat([confs, noise_est], dim=-1)
|
||||
weights, *_ = agent.get_action_and_value(state)
|
||||
fused = weights[:, 0:1]*tf + weights[:, 1:2]*af + weights[:, 2:3]*vf
|
||||
else:
|
||||
fused = (tf + af + vf) / 3.0
|
||||
logits = classifier(fused)
|
||||
all_preds.append(logits.argmax(-1).cpu())
|
||||
all_labels.append(labels.cpu())
|
||||
preds = torch.cat(all_preds).numpy()
|
||||
labels = torch.cat(all_labels).numpy()
|
||||
wf1 = float(f1_score(labels, preds, average="weighted", zero_division=0))
|
||||
acc = float(accuracy_score(labels, preds))
|
||||
return wf1, acc
|
||||
|
||||
|
||||
# ── Stage A: Supervised pretraining ──────────────────────────────────────
|
||||
|
||||
def train_stage_a(args, cfg, local_rank, rank, world_size):
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
rng = np.random.default_rng(42 + rank)
|
||||
|
||||
data_dir = os.path.join(PROJ, "data", args.dataset.lower())
|
||||
noise_root = os.path.join(PROJ, "data", f"{args.dataset.lower()}_noisy")
|
||||
|
||||
train_ds = MultimodalDataset(data_dir, "train", load_noisy=True,
|
||||
noise_root=noise_root)
|
||||
val_ds = MultimodalDataset(data_dir, "val")
|
||||
|
||||
train_loader = get_dataloader(train_ds, cfg["batch_size"], distributed=True)
|
||||
val_loader = get_dataloader(val_ds, cfg["batch_size"], shuffle=False,
|
||||
distributed=True, drop_last=False)
|
||||
|
||||
text_dim = train_ds.text.shape[1]
|
||||
audio_dim = train_ds.audio.shape[1]
|
||||
vision_dim = train_ds.vision.shape[1]
|
||||
num_classes = int(train_ds.labels.max()) + 1
|
||||
proj_dim = cfg.get("proj_dim", 1024)
|
||||
|
||||
encoder = MultimodalEncoder(text_dim, audio_dim, vision_dim, proj_dim).to(device)
|
||||
classifier = EmotionClassifier(proj_dim, num_classes,
|
||||
hidden=cfg.get("cls_hidden", 512)).to(device)
|
||||
encoder = DDP(encoder, device_ids=[local_rank])
|
||||
classifier = DDP(classifier, device_ids=[local_rank])
|
||||
|
||||
params = list(encoder.parameters()) + list(classifier.parameters())
|
||||
opt = torch.optim.AdamW(params, lr=cfg["lr"], weight_decay=cfg.get("wd", 1e-4))
|
||||
sched = torch.optim.lr_scheduler.CosineAnnealingLR(
|
||||
opt, T_max=cfg["epochs"], eta_min=1e-5)
|
||||
scaler = GradScaler()
|
||||
|
||||
conf_weight = cfg.get("conf_weight", 0.2)
|
||||
noise_prob = cfg.get("noise_prob", 0.4)
|
||||
best_wf1 = 0.0
|
||||
|
||||
for epoch in range(cfg["epochs"]):
|
||||
train_loader.sampler.set_epoch(epoch)
|
||||
encoder.train()
|
||||
classifier.train()
|
||||
ep_loss = ep_ce = ep_conf = 0.0
|
||||
|
||||
for batch in train_loader:
|
||||
text = batch["text"].to(device)
|
||||
audio = batch["audio"].to(device)
|
||||
vision = batch["vision"].to(device)
|
||||
labels = batch["labels"].to(device)
|
||||
B = text.size(0)
|
||||
|
||||
# Noise injection: randomly replace with noisy variant
|
||||
use_noise = (rng.random() < noise_prob) and bool(train_ds.variant_names)
|
||||
if use_noise:
|
||||
vname = rng.choice(train_ds.variant_names)
|
||||
v = train_ds.noisy_variants[vname]
|
||||
ni = rng.integers(0, len(train_ds), size=B)
|
||||
text, audio, vision, labels = _noisy_batch(train_ds, v, ni, device)
|
||||
|
||||
with autocast():
|
||||
tf, af, vf, confs = encoder(text, audio, vision)
|
||||
fused = (tf + af + vf) / 3.0
|
||||
logits = classifier(fused)
|
||||
ce_loss = F.cross_entropy(logits, labels)
|
||||
|
||||
# Confidence target: noisy modalities -> 0.1, clean -> 0.9
|
||||
if use_noise:
|
||||
c_tgt = _confidence_targets(vname, B, device)
|
||||
else:
|
||||
c_tgt = torch.full((B, 3), 0.9, device=device)
|
||||
conf_loss = F.binary_cross_entropy(confs, c_tgt)
|
||||
loss = ce_loss + conf_weight * conf_loss
|
||||
|
||||
opt.zero_grad(set_to_none=True)
|
||||
scaler.scale(loss).backward()
|
||||
scaler.unscale_(opt)
|
||||
nn.utils.clip_grad_norm_(params, 1.0)
|
||||
scaler.step(opt)
|
||||
scaler.update()
|
||||
|
||||
ep_loss += loss.item()
|
||||
ep_ce += ce_loss.item()
|
||||
ep_conf += conf_loss.item()
|
||||
|
||||
sched.step()
|
||||
val_wf1, val_acc = evaluate(encoder.module, classifier.module,
|
||||
val_loader, device)
|
||||
val_wf1 = all_reduce_mean(val_wf1, device)
|
||||
|
||||
if is_main(rank):
|
||||
n = len(train_loader)
|
||||
logging.info(
|
||||
f"[StageA] Epoch {epoch+1:3d}/{cfg['epochs']} | "
|
||||
f"loss={ep_loss/n:.4f} ce={ep_ce/n:.4f} conf={ep_conf/n:.4f} | "
|
||||
f"val_wf1={val_wf1:.4f} acc={val_acc:.4f}"
|
||||
)
|
||||
wandb.log({"A/loss": ep_loss/n, "A/ce": ep_ce/n,
|
||||
"A/conf": ep_conf/n, "A/val_wf1": val_wf1,
|
||||
"A/val_acc": val_acc, "epoch": epoch + 1})
|
||||
|
||||
if val_wf1 > best_wf1:
|
||||
best_wf1 = val_wf1
|
||||
save_ckpt({
|
||||
"epoch": epoch + 1,
|
||||
"encoder": encoder.module.state_dict(),
|
||||
"classifier": classifier.module.state_dict(),
|
||||
"val_wf1": val_wf1,
|
||||
"text_dim": text_dim, "audio_dim": audio_dim,
|
||||
"vision_dim": vision_dim, "num_classes": num_classes,
|
||||
"proj_dim": proj_dim, "cfg": cfg,
|
||||
}, os.path.join(args.output, "best.ckpt"))
|
||||
logging.info(f" -> New best WF1: {val_wf1:.4f}")
|
||||
|
||||
if is_main(rank):
|
||||
logging.info(f"Stage A done. Best val WF1: {best_wf1:.4f}")
|
||||
save_ckpt({
|
||||
"epoch": cfg["epochs"],
|
||||
"encoder": encoder.module.state_dict(),
|
||||
"classifier": classifier.module.state_dict(),
|
||||
"text_dim": text_dim, "audio_dim": audio_dim,
|
||||
"vision_dim": vision_dim, "num_classes": num_classes,
|
||||
"proj_dim": proj_dim, "cfg": cfg,
|
||||
}, os.path.join(args.output, "last.ckpt"))
|
||||
|
||||
dims = dict(text_dim=text_dim, audio_dim=audio_dim,
|
||||
vision_dim=vision_dim, num_classes=num_classes, proj_dim=proj_dim)
|
||||
return encoder.module, classifier.module, dims
|
||||
|
||||
|
||||
# ── Stage B: PPO training ─────────────────────────────────────────────────
|
||||
|
||||
def collect_rollout(encoder, classifier, agent, dataset, device, rollout_size, cfg, prev_weights):
|
||||
encoder.eval()
|
||||
classifier.eval()
|
||||
agent.eval()
|
||||
bs = cfg.get("batch_size", 128)
|
||||
nprob = cfg.get("noise_prob", 0.5)
|
||||
rng = np.random.default_rng()
|
||||
states, actions, log_probs, values, rewards = [], [], [], [], []
|
||||
collected = 0
|
||||
|
||||
with torch.no_grad():
|
||||
while collected < rollout_size:
|
||||
bsz = min(bs, rollout_size - collected)
|
||||
idx = rng.integers(0, len(dataset), size=bsz)
|
||||
|
||||
text = torch.from_numpy(dataset.text[idx]).to(device)
|
||||
audio = torch.from_numpy(dataset.audio[idx]).to(device)
|
||||
vision = torch.from_numpy(dataset.vision[idx]).to(device)
|
||||
labels = torch.from_numpy(dataset.labels[idx]).to(device)
|
||||
|
||||
if rng.random() < nprob and dataset.variant_names:
|
||||
vname = rng.choice(dataset.variant_names)
|
||||
v = dataset.noisy_variants[vname]
|
||||
text, audio, vision, labels = _noisy_batch(dataset, v, idx, device)
|
||||
|
||||
tf, af, vf, confs = encoder(text, audio, vision)
|
||||
noise_est = audio.std(dim=-1, keepdim=True).sigmoid()
|
||||
state = torch.cat([confs, noise_est], dim=-1) # (B, 4)
|
||||
|
||||
weights, log_p, value, _ = agent.get_action_and_value(state)
|
||||
fused = weights[:, 0:1]*tf + weights[:, 1:2]*af + weights[:, 2:3]*vf
|
||||
logits = classifier(fused)
|
||||
|
||||
rew, _ = compute_reward(
|
||||
logits, labels, confs, weights, prev_weights,
|
||||
alpha=cfg.get("reward_alpha", 1.0),
|
||||
beta =cfg.get("reward_beta", 0.3),
|
||||
gamma=cfg.get("reward_gamma", 0.1),
|
||||
)
|
||||
|
||||
states.append(state)
|
||||
actions.append(weights)
|
||||
log_probs.append(log_p)
|
||||
values.append(value.squeeze(-1))
|
||||
rewards.append(rew)
|
||||
collected += bsz
|
||||
|
||||
states = torch.cat(states)
|
||||
actions = torch.cat(actions)
|
||||
log_probs = torch.cat(log_probs)
|
||||
values = torch.cat(values)
|
||||
rewards = torch.cat(rewards)
|
||||
|
||||
rewards = (rewards - rewards.mean()) / (rewards.std() + 1e-8)
|
||||
advantages = rewards - values.detach().cpu()
|
||||
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
|
||||
|
||||
return dict(states=states, actions=actions, log_probs=log_probs,
|
||||
values=values, rewards=rewards, advantages=advantages,
|
||||
mean_weights=actions.mean(0))
|
||||
|
||||
|
||||
def ppo_update(agent, opt, rollout, cfg, device, scaler):
|
||||
eps = cfg.get("ppo_clip", 0.2)
|
||||
ppo_ep = cfg.get("ppo_epochs_per_update", 4)
|
||||
mb_size = cfg.get("ppo_mini_batch", 256)
|
||||
v_coef = cfg.get("value_coef", 0.5)
|
||||
ent_coef = cfg.get("entropy_coef", 0.01)
|
||||
|
||||
states = rollout["states"].to(device)
|
||||
actions = rollout["actions"].to(device)
|
||||
old_lp = rollout["log_probs"].to(device)
|
||||
adv = rollout["advantages"].to(device)
|
||||
ret = rollout["rewards"].to(device)
|
||||
n = states.size(0)
|
||||
|
||||
total_pl = total_vl = total_ent = cnt = 0.0
|
||||
agent.train()
|
||||
|
||||
for _ in range(ppo_ep):
|
||||
perm = torch.randperm(n, device=device)
|
||||
for start in range(0, n, mb_size):
|
||||
idx = perm[start:start + mb_size]
|
||||
s = states[idx]; a = actions[idx]
|
||||
olp = old_lp[idx]; ad = adv[idx]; r = ret[idx]
|
||||
|
||||
with autocast():
|
||||
new_lp, val, ent = agent.evaluate(s, a)
|
||||
val = val.squeeze(-1)
|
||||
ratio = (new_lp - olp).exp()
|
||||
p_loss = -torch.min(ratio * ad,
|
||||
torch.clamp(ratio, 1-eps, 1+eps) * ad).mean()
|
||||
v_loss = F.mse_loss(val, r)
|
||||
e_loss = -ent.mean()
|
||||
loss = p_loss + v_coef * v_loss + ent_coef * e_loss
|
||||
|
||||
opt.zero_grad(set_to_none=True)
|
||||
scaler.scale(loss).backward()
|
||||
scaler.unscale_(opt)
|
||||
nn.utils.clip_grad_norm_(agent.parameters(), 0.5)
|
||||
scaler.step(opt)
|
||||
scaler.update()
|
||||
|
||||
total_pl += p_loss.item()
|
||||
total_vl += v_loss.item()
|
||||
total_ent += ent.mean().item()
|
||||
cnt += 1
|
||||
|
||||
return dict(p_loss=total_pl/cnt, v_loss=total_vl/cnt, entropy=total_ent/cnt)
|
||||
|
||||
|
||||
def train_stage_b(args, cfg, encoder, classifier, dims, local_rank, rank, world_size):
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
|
||||
data_dir = os.path.join(PROJ, "data", args.dataset.lower())
|
||||
noise_root = os.path.join(PROJ, "data", f"{args.dataset.lower()}_noisy")
|
||||
train_ds = MultimodalDataset(data_dir, "train", load_noisy=True,
|
||||
noise_root=noise_root)
|
||||
val_ds = MultimodalDataset(data_dir, "val")
|
||||
val_loader = get_dataloader(val_ds, cfg.get("batch_size", 128),
|
||||
shuffle=False, distributed=True, drop_last=False)
|
||||
|
||||
# Freeze encoder (projectors + confidence estimators)
|
||||
for p in encoder.parameters():
|
||||
p.requires_grad_(False)
|
||||
encoder.to(device).eval()
|
||||
|
||||
# Classifier: keep trainable (supervised component)
|
||||
classifier.to(device)
|
||||
cls_ddp = DDP(classifier, device_ids=[local_rank])
|
||||
opt_cls = torch.optim.AdamW(classifier.parameters(),
|
||||
lr=cfg.get("cls_lr", 5e-5), weight_decay=1e-4)
|
||||
|
||||
# RL agent
|
||||
agent = ModalFusionAgent(state_dim=4,
|
||||
hidden=cfg.get("agent_hidden", 128)).to(device)
|
||||
agent = DDP(agent, device_ids=[local_rank])
|
||||
opt_agent = torch.optim.Adam(agent.parameters(), lr=cfg.get("rl_lr", 3e-4))
|
||||
|
||||
scaler = GradScaler()
|
||||
|
||||
rollout_size = cfg.get("rollout_steps", 512)
|
||||
n_updates = cfg.get("n_ppo_updates", 500)
|
||||
eval_every = cfg.get("eval_every", 10)
|
||||
best_wf1 = 0.0
|
||||
prev_weights = None
|
||||
|
||||
for upd in range(n_updates):
|
||||
# Rollout collection
|
||||
rollout = collect_rollout(
|
||||
encoder, classifier, agent.module,
|
||||
train_ds, device, rollout_size, cfg, prev_weights,
|
||||
)
|
||||
prev_weights = rollout["mean_weights"].to(device)
|
||||
|
||||
# PPO update
|
||||
ppo_info = ppo_update(agent, opt_agent, rollout, cfg, device, scaler)
|
||||
|
||||
# Lightweight supervised classifier refresh (one mini-batch every 2 updates)
|
||||
if upd % 2 == 0:
|
||||
idx = np.random.randint(0, len(train_ds), cfg.get("batch_size", 128))
|
||||
text = torch.from_numpy(train_ds.text[idx]).to(device)
|
||||
audio = torch.from_numpy(train_ds.audio[idx]).to(device)
|
||||
vision = torch.from_numpy(train_ds.vision[idx]).to(device)
|
||||
labels = torch.from_numpy(train_ds.labels[idx]).to(device)
|
||||
with torch.no_grad():
|
||||
tf, af, vf, confs = encoder(text, audio, vision)
|
||||
noise_est = audio.std(dim=-1, keepdim=True).sigmoid()
|
||||
state = torch.cat([confs, noise_est], dim=-1)
|
||||
weights, *_ = agent.module.get_action_and_value(state)
|
||||
fused = weights[:, 0:1]*tf + weights[:, 1:2]*af + weights[:, 2:3]*vf
|
||||
with autocast():
|
||||
logits = cls_ddp(fused)
|
||||
loss = F.cross_entropy(logits, labels)
|
||||
opt_cls.zero_grad(set_to_none=True)
|
||||
scaler.scale(loss).backward()
|
||||
scaler.step(opt_cls)
|
||||
scaler.update()
|
||||
|
||||
# Evaluate
|
||||
if upd % eval_every == 0:
|
||||
val_wf1, val_acc = evaluate(encoder, classifier, val_loader, device,
|
||||
agent=agent.module)
|
||||
val_wf1 = all_reduce_mean(val_wf1, device)
|
||||
|
||||
if is_main(rank):
|
||||
mean_rew = rollout["rewards"].mean().item()
|
||||
logging.info(
|
||||
f"[StageB] PPO {upd:4d}/{n_updates} | "
|
||||
f"rew={mean_rew:.4f} p={ppo_info['p_loss']:.4f} "
|
||||
f"v={ppo_info['v_loss']:.4f} ent={ppo_info['entropy']:.4f} | "
|
||||
f"val_wf1={val_wf1:.4f}"
|
||||
)
|
||||
wandb.log({
|
||||
"B/reward": mean_rew,
|
||||
"B/p_loss": ppo_info["p_loss"],
|
||||
"B/v_loss": ppo_info["v_loss"],
|
||||
"B/entropy": ppo_info["entropy"],
|
||||
"B/val_wf1": val_wf1,
|
||||
"ppo_update": upd,
|
||||
})
|
||||
|
||||
if val_wf1 > best_wf1:
|
||||
best_wf1 = val_wf1
|
||||
save_ckpt({
|
||||
"update": upd,
|
||||
"encoder": encoder.state_dict(),
|
||||
"classifier": classifier.state_dict(),
|
||||
"agent": agent.module.state_dict(),
|
||||
"val_wf1": val_wf1,
|
||||
**dims,
|
||||
}, os.path.join(args.output, "best.ckpt"))
|
||||
logging.info(f" -> New best WF1: {val_wf1:.4f}")
|
||||
|
||||
if is_main(rank):
|
||||
logging.info(f"Stage B done. Best val WF1: {best_wf1:.4f}")
|
||||
|
||||
|
||||
# ── Main ──────────────────────────────────────────────────────────────────
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--stage", required=True, choices=["supervised", "rl"])
|
||||
p.add_argument("--dataset", default="IEMOCAP")
|
||||
p.add_argument("--config", required=True)
|
||||
p.add_argument("--output", required=True)
|
||||
p.add_argument("--checkpoint", default=None)
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
local_rank, rank, world_size = setup_ddp()
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
|
||||
os.makedirs(args.output, exist_ok=True)
|
||||
log_dir = os.path.join(PROJ, "outputs", "logs")
|
||||
os.makedirs(log_dir, exist_ok=True)
|
||||
|
||||
stage_tag = "stageA" if args.stage == "supervised" else "stageB"
|
||||
handlers = [logging.StreamHandler()]
|
||||
if is_main(rank):
|
||||
handlers.append(logging.FileHandler(
|
||||
os.path.join(log_dir, f"{stage_tag}.log"), mode="a"))
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format=f"[rank{rank}] %(asctime)s %(message)s",
|
||||
handlers=handlers,
|
||||
)
|
||||
|
||||
with open(os.path.join(PROJ, args.config)) as f:
|
||||
cfg = yaml.safe_load(f)
|
||||
|
||||
if is_main(rank):
|
||||
os.environ.setdefault("WANDB_MODE", "offline")
|
||||
wandb.init(
|
||||
project="multimodal_affect",
|
||||
name=f"d1_{args.stage}_{args.dataset}_{time.strftime('%m%d_%H%M')}",
|
||||
config={**cfg, "stage": args.stage, "dataset": args.dataset},
|
||||
dir=os.path.join(PROJ, "outputs"),
|
||||
)
|
||||
|
||||
if args.stage == "supervised":
|
||||
train_stage_a(args, cfg, local_rank, rank, world_size)
|
||||
|
||||
elif args.stage == "rl":
|
||||
if not args.checkpoint:
|
||||
raise ValueError("--checkpoint required for --stage rl")
|
||||
ckpt = torch.load(args.checkpoint, map_location=device)
|
||||
|
||||
text_dim = ckpt["text_dim"]
|
||||
audio_dim = ckpt["audio_dim"]
|
||||
vision_dim = ckpt["vision_dim"]
|
||||
num_classes = ckpt["num_classes"]
|
||||
proj_dim = ckpt.get("proj_dim", 1024)
|
||||
dims = dict(text_dim=text_dim, audio_dim=audio_dim,
|
||||
vision_dim=vision_dim, num_classes=num_classes,
|
||||
proj_dim=proj_dim)
|
||||
|
||||
encoder = MultimodalEncoder(text_dim, audio_dim, vision_dim, proj_dim)
|
||||
classifier = EmotionClassifier(proj_dim, num_classes)
|
||||
encoder.load_state_dict(ckpt["encoder"])
|
||||
classifier.load_state_dict(ckpt["classifier"])
|
||||
|
||||
if is_main(rank):
|
||||
logging.info(
|
||||
f"Loaded Stage A ckpt from {args.checkpoint} "
|
||||
f"(val_wf1={ckpt.get('val_wf1', 0.0):.4f})"
|
||||
)
|
||||
train_stage_b(args, cfg, encoder, classifier, dims,
|
||||
local_rank, rank, world_size)
|
||||
|
||||
if is_main(rank):
|
||||
wandb.finish()
|
||||
cleanup()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
'''
|
||||
|
||||
# ─── configs/d1/stage_a.yaml ──────────────────────────────────────────────
|
||||
STAGE_A_YAML = '''\
|
||||
# Stage A: Supervised pretraining
|
||||
# Trains projection MLPs + confidence estimators + classifier
|
||||
# with noise injection to teach confidence estimation
|
||||
|
||||
epochs: 50
|
||||
batch_size: 128
|
||||
lr: 2.0e-4
|
||||
wd: 1.0e-4
|
||||
proj_dim: 1024
|
||||
cls_hidden: 512
|
||||
conf_weight: 0.2 # BCE loss weight for confidence estimators
|
||||
noise_prob: 0.4 # probability of injecting noisy batch
|
||||
'''
|
||||
|
||||
# ─── configs/d1/stage_b.yaml ──────────────────────────────────────────────
|
||||
STAGE_B_YAML = '''\
|
||||
# Stage B: PPO-based adaptive fusion weight learning
|
||||
# Encoder (projectors + confidence estimators) frozen from Stage A
|
||||
# RL agent learns noise-adaptive fusion weights via PPO
|
||||
|
||||
batch_size: 128
|
||||
proj_dim: 1024
|
||||
|
||||
# PPO
|
||||
rollout_steps: 512 # experiences collected per PPO update
|
||||
n_ppo_updates: 500 # total PPO update iterations
|
||||
ppo_clip: 0.2
|
||||
ppo_epochs_per_update: 4
|
||||
ppo_mini_batch: 256
|
||||
value_coef: 0.5
|
||||
entropy_coef: 0.01
|
||||
rl_lr: 3.0e-4
|
||||
cls_lr: 5.0e-5
|
||||
|
||||
# Reward coefficients (R = alpha*(-CE) + beta*Consistency - gamma*Instability)
|
||||
reward_alpha: 1.0
|
||||
reward_beta: 0.3
|
||||
reward_gamma: 0.1
|
||||
|
||||
# RL agent architecture
|
||||
agent_hidden: 128
|
||||
|
||||
# Noise injection during rollout collection
|
||||
noise_prob: 0.5
|
||||
|
||||
eval_every: 10 # evaluate every N PPO updates
|
||||
'''
|
||||
|
||||
# Upload
|
||||
uploads = {
|
||||
f"{PROJ}/scripts/train/train_d1.py": TRAIN_D1,
|
||||
f"{PROJ}/configs/d1/stage_a.yaml": STAGE_A_YAML,
|
||||
f"{PROJ}/configs/d1/stage_b.yaml": STAGE_B_YAML,
|
||||
}
|
||||
|
||||
for path, content in uploads.items():
|
||||
with sftp.open(path, 'w') as f:
|
||||
f.write(content)
|
||||
print(f" uploaded: {path.split('multimodal_affect/')[-1]}")
|
||||
|
||||
sftp.close()
|
||||
client.close()
|
||||
print("\nAll training files uploaded.")
|
||||
Reference in New Issue
Block a user