| 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(save_dir, rank): |
| os.makedirs(save_dir, exist_ok=True) |
| logger = logging.getLogger("TrainLogger") |
| if rank == 0: |
| logger.setLevel(logging.INFO) |
| formatter = logging.Formatter('%(asctime)s - %(message)s') |
| file_handler = logging.FileHandler(os.path.join(save_dir, f"train_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log")) |
| file_handler.setFormatter(formatter) |
| logger.addHandler(file_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() |
|
|
| logger = setup_logger(args.save_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) |
| |
| |
| 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=150) |
| |
| for batch in 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[:, :t_idx, :, :] |
| target = latents[:, t_idx, :, :] |
| |
| 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)) |
| |
| total_loss = criterion(next_t_init, target) + criterion(next_t_refined, target) |
| total_loss.backward() |
| optimizer.step() |
| |
| if rank == 0: |
| pbar.set_postfix({"Loss": f"{total_loss.item():.3f}"}) |
| |
| |
| if rank == 0 and (epoch + 1) % args.save_every_k_epochs == 0: |
| ckpt_path = os.path.join(args.save_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) |
