| 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): |
| |
| 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) |
| |
| |
| 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) |
| |
| |
| 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=4) |
| bar_format = '{l_bar}{bar:30}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]' |
| from itertools import islice |
| |
| skipped_dl = islice(dataloader, 17010, None) |
| |
| |
| 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) |
| |
| 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}"}) |
| |
| |
| 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=2048) |
| 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=1e-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) |
