| import os |
| import torch |
| import torch.nn as nn |
| import argparse |
| import torch.optim as optim |
| import logging |
| from datetime import datetime |
| from tqdm import tqdm |
|
|
| |
| 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): |
| """两层线性层 + 残差连接：将 512 维投影至 64 维""" |
| 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): |
| os.makedirs(save_dir, exist_ok=True) |
| time_str = datetime.now().strftime("%Y%m%d_%H%M%S") |
| log_file = os.path.join(save_dir, f"train_{time_str}.log") |
| |
| logger = logging.getLogger("TrainLogger") |
| logger.setLevel(logging.INFO) |
| |
| formatter = logging.Formatter('%(asctime)s - %(message)s') |
| |
| 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) |
| |
| 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): |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| logger = setup_logger(args.save_dir) |
| |
| |
| 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) |
| |
| logger.info("--- 参数统计 ---") |
| count_parameters(temporal_model, "TemporalModel", logger) |
| count_parameters(spatial_model, "SpatialModel", logger) |
| count_parameters(text_proj, "TextProjection", logger) |
| |
| criterion = TokenL2Loss() |
| optimizer = optim.Adam( |
| list(temporal_model.parameters()) + list(spatial_model.parameters()) + list(text_proj.parameters()), |
| lr=args.lr |
| ) |
| |
| 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}]' |
| |
| |
| for epoch in range(args.epochs): |
| temporal_model.train() |
| spatial_model.train() |
| text_proj.train() |
| |
| pbar = tqdm(dataloader, desc=f"Epoch {epoch+1:03d}/{args.epochs:03d}", |
| ncols=300, bar_format=bar_format) |
| |
| for batch in pbar: |
| optimizer.zero_grad() |
| |
| vertices = batch['vertices'].to(device) |
| faces = batch['faces'].to(device) |
|
|
| max_v_idx = vertices.shape[2] - 1 |
| if faces.max() > max_v_idx: |
| pbar.write(f"\n[!] 发现非法数据，索引 {faces.max()} 超出范围，跳过Batch") |
| continue |
| |
| with torch.no_grad(): |
| B, T, V, C = vertices.shape |
| latents = [] |
| for t in range(T): |
| l = vae.encode_mesh(vertices[:, t].squeeze(0), faces.squeeze(0)) |
| latents.append(l) |
| latents = torch.stack(latents, dim=1) |
| |
| 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)) |
| |
| loss_t = criterion(next_t_init, target) |
| loss_s = criterion(next_t_refined, target) |
| total_loss = loss_t + loss_s |
| |
| total_loss.backward() |
| optimizer.step() |
| |
| pbar.set_postfix({ |
| "Loss": f"{total_loss.item():.3f}", |
| "Temp": f"{loss_t.item():.3f}", |
| "Spat": f"{loss_s.item():.3f}" |
| }) |
| |
| logger.info(f"Epoch {epoch+1}/{args.epochs} 完成 - Total Loss: {total_loss.item():.4f}") |
| |
| if (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.state_dict(), |
| 'spatial': spatial_model.state_dict(), |
| 'proj': text_proj.state_dict() |
| }, ckpt_path) |
| logger.info(f"模型已保存至: {ckpt_path}") |
|
|
| 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) |
