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) 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=6) bar_format = '{l_bar}{bar:30}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]' # 3. 训练循环 global_step = 0 for epoch in range(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) 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 + loss_s 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: # 每2000个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(), '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: {loss_t.item():.4f} | Spat: {loss_s.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.item():.3f}", "Spat": f"{loss_s.item():.3f}"}) # 4. 主进程保存模型到专属子文件夹 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() }, 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") args = parser.parse_args() train(args)