import os import torch import torch.nn as nn import argparse import torch.optim as optim import logging import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from datetime import datetime from tqdm import tqdm import numpy as np # 假设你的组件路径正确 from models.auto import TemporalTrajectoryPredictor, SpatialGraphRefiner from models.vae import TripoSGVaeWrapper from models.clip import FrozenCLIPTextEncoder from utils.autodataloader import get_mesh_dataloader from utils.loss import TokenL2Loss class TextProjectionLayer(nn.Module): def __init__(self, in_dim=512, out_dim=64): super().__init__() self.net = nn.Sequential(nn.Linear(in_dim, in_dim), nn.ReLU(), nn.Linear(in_dim, out_dim)) self.residual = nn.Linear(in_dim, out_dim) if in_dim != out_dim else nn.Identity() def forward(self, x): return self.net(x) + self.residual(x) def setup_logger(run_dir, rank): logger = logging.getLogger("TrainLogger") if logger.hasHandlers(): logger.handlers.clear() if rank == 0: logger.setLevel(logging.INFO) formatter = logging.Formatter('%(asctime)s - %(message)s') log_file = os.path.join(run_dir, "train.log") file_handler = logging.FileHandler(log_file) file_handler.setFormatter(formatter) logger.addHandler(file_handler) console_handler = logging.StreamHandler() console_handler.setFormatter(formatter) logger.addHandler(console_handler) else: logger.setLevel(logging.WARNING) return logger def count_parameters(model, name, logger): trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) logger.info(f"[{name}] 可训练参数量: {trainable:,}") return trainable def train(args): # 1. 初始化分布式环境 local_rank = int(os.environ["LOCAL_RANK"]) torch.cuda.set_device(local_rank) device = torch.device("cuda", local_rank) dist.init_process_group(backend="nccl", device_id=device) rank = dist.get_rank() # 创建本次运行的专属文件夹 run_name = datetime.now().strftime("%Y%m%d_%H%M%S") run_dir = os.path.join(args.save_dir, run_name) if rank == 0: os.makedirs(run_dir, exist_ok=True) dist.barrier() logger = setup_logger(run_dir, rank) # 2. 初始化模型 vae = TripoSGVaeWrapper(device=device) vae.eval() for param in vae.parameters(): param.requires_grad = False text_encoder = FrozenCLIPTextEncoder(args.clip_model_path, device=device) text_proj = TextProjectionLayer(in_dim=512, out_dim=args.token_dim).to(device) temporal_model = TemporalTrajectoryPredictor(num_tokens=args.num_tokens, token_dim=args.token_dim, max_seq_len=16).to(device) spatial_model = SpatialGraphRefiner(num_tokens=args.num_tokens, token_dim=args.token_dim, text_dim=args.token_dim).to(device) # 包装 DDP text_proj = DDP(text_proj, device_ids=[local_rank]) temporal_model = DDP(temporal_model, device_ids=[local_rank]) spatial_model = DDP(spatial_model, device_ids=[local_rank]) optimizer = optim.Adam( list(temporal_model.parameters()) + list(spatial_model.parameters()) + list(text_proj.parameters()), lr=args.lr ) logger.info("--- 参数统计 ---") count_parameters(temporal_model, "TemporalModel", logger) count_parameters(spatial_model, "SpatialModel", logger) count_parameters(text_proj, "TextProjection", logger) # 恢复训练的状态初始化 start_epoch = 0 global_step = 0 # 判断是否启用自动恢复 if args.auto_resume: if args.last_ckpt and os.path.exists(args.last_ckpt): logger.info(f"正在从检查点恢复训练: {args.last_ckpt}") # map_location="cpu" 避免多卡同时加载引起的显存瞬间突刺 checkpoint = torch.load(args.last_ckpt, map_location="cpu") # 恢复模型权重（因为包装了 DDP，所以需要给 module 赋值） temporal_model.module.load_state_dict(checkpoint['temporal']) spatial_model.module.load_state_dict(checkpoint['spatial']) text_proj.module.load_state_dict(checkpoint['proj']) # 恢复优化器状态 if 'optimizer' in checkpoint: optimizer.load_state_dict(checkpoint['optimizer']) # 确保优化器内部的张量和当前设备一致 for state in optimizer.state.values(): for k, v in state.items(): if isinstance(v, torch.Tensor): state[k] = v.to(device) # 恢复训练计数状态（针对 step 保存的完整检查点） if 'epoch' in checkpoint: start_epoch = checkpoint['epoch'] + 1 # 从下一轮开始 if 'step' in checkpoint: global_step = checkpoint['step'] logger.info(f"恢复成功！将从 Epoch {start_epoch+1}，Global Step {global_step} 继续训练。") else: logger.warning(f"指定了 auto_resume=True，但未找到有效检查点路径: {args.last_ckpt}。模型将从头启动训练。") criterion = TokenL2Loss() data_dir = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/data/shards/" dataloader = get_mesh_dataloader(data_dir=data_dir, batch_size=args.batch_size, num_workers=12) bar_format = '{l_bar}{bar:30}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]' # 3. 训练循环（将 range(args.epochs) 改为从 start_epoch 开始） for epoch in range(start_epoch, args.epochs): temporal_model.train() spatial_model.train() text_proj.train() pbar = tqdm(dataloader, desc=f"Epoch {epoch+1:03d}", disable=(rank != 0), ncols=1200, bar_format=bar_format) for batch_idx, batch in enumerate(pbar): optimizer.zero_grad() vertices, faces = batch['vertices'], batch['faces'] with torch.no_grad(): latents_batch = [] for i in range(len(vertices)): sample_latents = [] for t in range(vertices[i].shape[0]): v_t = vertices[i][t].detach().clone().to(device=device, dtype=torch.float32) f_t = faces[i].detach().clone().to(device=device, dtype=torch.float32) l = vae.encode_mesh(v_t, f_t) if isinstance(l, np.ndarray): l = torch.from_numpy(l).to(device) sample_latents.append(l.squeeze(0)) latents_batch.append(torch.stack(sample_latents, dim=0)) latents = torch.stack(latents_batch, dim=0) text_embed = text_proj(text_encoder(list(batch['caption']))).unsqueeze(1) history = latents[:, :-1, :, :] target = latents[:, 1:, :, :] next_t_init = temporal_model(history, text_embed) if False: print("target mean/std:", target.mean().item(), target.std().item()) print("pred mean/std:", next_t_init.mean().item(), next_t_init.std().item()) print("target norm:", target.norm(dim=-1).mean().item()) print("pred norm:", next_t_init.norm(dim=-1).mean().item()) print("input norm:", history.norm(dim=-1).mean().item()) next_t_refined = spatial_model(next_t_init, text_embed, torch.zeros((len(batch['vertices']), args.num_tokens, args.num_tokens)).to(device)) loss_t = criterion(next_t_init, target) loss_s = criterion(next_t_refined, target) total_loss = loss_t['total'] + loss_s['total'] if torch.isnan(total_loss): print(f"\n[!] NaN detected at Epoch {epoch+1}, Batch {batch_idx+1}") raise ValueError("Training stopped: NaN loss detected.") total_loss.backward() torch.nn.utils.clip_grad_norm_(list(temporal_model.parameters()) + list(spatial_model.parameters()) + list(text_proj.parameters()), max_norm=1.0) optimizer.step() global_step += 1 if rank == 0: # 每500个batch保存一次 if global_step % 500 == 0: ckpt_path = os.path.join(run_dir, f"model_step_{global_step}.pth") torch.save({ 'temporal': temporal_model.module.state_dict(), 'spatial': spatial_model.module.state_dict(), 'proj': text_proj.module.state_dict(), 'optimizer': optimizer.state_dict(), # 存入优化器状态以便精准恢复 'epoch': epoch, 'step': global_step }, ckpt_path) logger.info(f"检查点已保存 (Step {global_step}): {ckpt_path}") if (batch_idx + 1) % 10 == 0: log_msg = f"Epoch {epoch+1:03d} | Batch {batch_idx+1:04d} | Total: {total_loss.item():.4f} | Temp L1: {loss_t['l1'].item():.4f} | Temp L2: {loss_t['l2'].item():.4f} | Spat L1: {loss_s['l1'].item():.4f} | Spat L2: {loss_s['l2'].item():.4f}" logger.info(log_msg) for handler in logger.handlers: handler.flush() pbar.set_postfix({"Loss": f"{total_loss.item():.3f}", "Temp": f"{loss_t['total'].item():.3f}", "Spat": f"{loss_s['total'].item():.3f}"}) # 4. 主进程按 Epoch 保存模型 if rank == 0 and (epoch + 1) % args.save_every_k_epochs == 0: ckpt_path = os.path.join(run_dir, f"model_epoch_{epoch+1}.pth") torch.save({ 'temporal': temporal_model.module.state_dict(), 'spatial': spatial_model.module.state_dict(), 'proj': text_proj.module.state_dict(), 'optimizer': optimizer.state_dict(), # Epoch 级别检查点也保存优化器状态 'epoch': epoch, 'step': global_step }, ckpt_path) logger.info(f"模型已保存: {ckpt_path}") dist.destroy_process_group() if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--num_tokens", type=int, default=512) parser.add_argument("--token_dim", type=int, default=64) parser.add_argument("--batch_size", type=int, default=4) parser.add_argument("--lr", type=float, default=5e-4) parser.add_argument("--epochs", type=int, default=100) parser.add_argument("--save_every_k_epochs", type=int, default=10) parser.add_argument("--save_dir", type=str, default="./checkpoints") parser.add_argument("--clip_model_path", type=str, default="/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/pretrain/clip-vit-base-patch32") # 新增的控制恢复训练参数 parser.add_argument("--auto_resume", action="store_true", help="是否自动从检查点恢复训练") parser.add_argument("--last_ckpt", type=str, default="", help="需要恢复训练的权重文件路径 (.pth)") args = parser.parse_args() train(args)