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 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") 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 ) criterion = TokenL2Loss() dataloader = get_mesh_dataloader(batch_size=args.batch_size, num_workers=80) bar_format = '{l_bar}{bar:30}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]' from itertools import islice # skipped_dl = islice(dataloader, 3260, None) skipped_dl = islice(dataloader, 0, None) # 3. 训练循环 for epoch in range(args.epochs): temporal_model.train() spatial_model.train() text_proj.train() pbar = tqdm(skipped_dl, 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 = torch.tensor(vertices[i][t], dtype=torch.float32, device=device) f_t = torch.tensor(faces[i], dtype=torch.long, device=device) 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) # print(f"latents.shape: {latents.shape}") text_embed = text_proj(text_encoder(list(batch['caption']))).unsqueeze(1) t_idx = torch.randint(1, 16, (1,)).item() history = latents[:, :-1, :, :] target = latents[:, 1:, :, :] next_t_init = temporal_model(history) next_t_refined = spatial_model(next_t_init, text_embed, torch.zeros((args.batch_size, 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}") print(f"Total Loss: {total_loss.item()}, Temp Loss: {loss_t.item()}, Spat Loss: {loss_s.item()}") print(f"History shape: {history.shape}, Data: {history}") print(f"Target shape: {target.shape}, Data: {target}") print(f"Next_t_init: {next_t_init}") print(f"Text Embed: {text_embed}") 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() if rank == 0: 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=4e-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)