AR / train_v6.py

Upload folder using huggingface_hub

394918c verified 4 months ago

11.4 kB

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)
    
    # 恢复训练的状态初始化
    start_epoch = 0
    global_step = 0

    # 判断是否启用自动恢复
    if args.auto_resume:
        if args.last_ckpt and os.path.exists(args.last_ckpt):
            logger.info(f"正在从检查点恢复训练: {args.last_ckpt}")
            # map_location="cpu" 避免多卡同时加载引起的显存瞬间突刺
            checkpoint = torch.load(args.last_ckpt, map_location="cpu")
            
            # 恢复模型权重（因为包装了 DDP，所以需要给 module 赋值）
            temporal_model.module.load_state_dict(checkpoint['temporal'])
            spatial_model.module.load_state_dict(checkpoint['spatial'])
            text_proj.module.load_state_dict(checkpoint['proj'])
            
            # 恢复优化器状态
            if 'optimizer' in checkpoint:
                optimizer.load_state_dict(checkpoint['optimizer'])
                # 确保优化器内部的张量和当前设备一致
                for state in optimizer.state.values():
                    for k, v in state.items():
                        if isinstance(v, torch.Tensor):
                            state[k] = v.to(device)
            
            # 恢复训练计数状态（针对 step 保存的完整检查点）
            if 'epoch' in checkpoint:
                start_epoch = checkpoint['epoch'] + 1  # 从下一轮开始
            if 'step' in checkpoint:
                global_step = checkpoint['step']
                
            logger.info(f"恢复成功！将从 Epoch {start_epoch+1}，Global Step {global_step} 继续训练。")
        else:
            logger.warning(f"指定了 auto_resume=True，但未找到有效检查点路径: {args.last_ckpt}。模型将从头启动训练。")

    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=12)
    
    bar_format = '{l_bar}{bar:30}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
    
    # 3. 训练循环（将 range(args.epochs) 改为从 start_epoch 开始）
    for epoch in range(start_epoch, 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, text_embed)
            if False:
                print("target mean/std:", target.mean().item(), target.std().item())
                print("pred mean/std:", next_t_init.mean().item(), next_t_init.std().item())
                print("target norm:", target.norm(dim=-1).mean().item())
                print("pred norm:", next_t_init.norm(dim=-1).mean().item())
                print("input norm:", history.norm(dim=-1).mean().item())
            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['total'] + loss_s['total']
            
            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:
                # 每500个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(),
                        'optimizer': optimizer.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 L1: {loss_t['l1'].item():.4f} | Temp L2: {loss_t['l2'].item():.4f} | Spat L1: {loss_s['l1'].item():.4f} | Spat L2: {loss_s['l2'].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['total'].item():.3f}", "Spat": f"{loss_s['total'].item():.3f}"})
        
        # 4. 主进程按 Epoch 保存模型
        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(),
                'optimizer': optimizer.state_dict(),  # Epoch 级别检查点也保存优化器状态
                'epoch': epoch,
                'step': global_step
            }, 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")
    
    # 新增的控制恢复训练参数
    parser.add_argument("--auto_resume", action="store_true", help="是否自动从检查点恢复训练")
    parser.add_argument("--last_ckpt", type=str, default="", help="需要恢复训练的权重文件路径 (.pth)")
    
    args = parser.parse_args()
    train(args)

	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)

	# 恢复训练的状态初始化
	start_epoch = 0
	global_step = 0

	# 判断是否启用自动恢复
	if args.auto_resume:
	if args.last_ckpt and os.path.exists(args.last_ckpt):
	logger.info(f"正在从检查点恢复训练: {args.last_ckpt}")
	# map_location="cpu" 避免多卡同时加载引起的显存瞬间突刺
	checkpoint = torch.load(args.last_ckpt, map_location="cpu")

	# 恢复模型权重（因为包装了 DDP，所以需要给 module 赋值）
	temporal_model.module.load_state_dict(checkpoint['temporal'])
	spatial_model.module.load_state_dict(checkpoint['spatial'])
	text_proj.module.load_state_dict(checkpoint['proj'])

	# 恢复优化器状态
	if 'optimizer' in checkpoint:
	optimizer.load_state_dict(checkpoint['optimizer'])
	# 确保优化器内部的张量和当前设备一致
	for state in optimizer.state.values():
	for k, v in state.items():
	if isinstance(v, torch.Tensor):
	state[k] = v.to(device)

	# 恢复训练计数状态（针对 step 保存的完整检查点）
	if 'epoch' in checkpoint:
	start_epoch = checkpoint['epoch'] + 1 # 从下一轮开始
	if 'step' in checkpoint:
	global_step = checkpoint['step']

	logger.info(f"恢复成功！将从 Epoch {start_epoch+1}，Global Step {global_step} 继续训练。")
	else:
	logger.warning(f"指定了 auto_resume=True，但未找到有效检查点路径: {args.last_ckpt}。模型将从头启动训练。")

	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=12)

	bar_format = '{l_bar}{bar:30}\| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'

	# 3. 训练循环（将 range(args.epochs) 改为从 start_epoch 开始）
	for epoch in range(start_epoch, 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, text_embed)
	if False:
	print("target mean/std:", target.mean().item(), target.std().item())
	print("pred mean/std:", next_t_init.mean().item(), next_t_init.std().item())
	print("target norm:", target.norm(dim=-1).mean().item())
	print("pred norm:", next_t_init.norm(dim=-1).mean().item())
	print("input norm:", history.norm(dim=-1).mean().item())
	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['total'] + loss_s['total']

	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:
	# 每500个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(),
	'optimizer': optimizer.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 L1: {loss_t['l1'].item():.4f} \| Temp L2: {loss_t['l2'].item():.4f} \| Spat L1: {loss_s['l1'].item():.4f} \| Spat L2: {loss_s['l2'].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['total'].item():.3f}", "Spat": f"{loss_s['total'].item():.3f}"})

	# 4. 主进程按 Epoch 保存模型
	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(),
	'optimizer': optimizer.state_dict(), # Epoch 级别检查点也保存优化器状态
	'epoch': epoch,
	'step': global_step
	}, 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")

	# 新增的控制恢复训练参数
	parser.add_argument("--auto_resume", action="store_true", help="是否自动从检查点恢复训练")
	parser.add_argument("--last_ckpt", type=str, default="", help="需要恢复训练的权重文件路径 (.pth)")

	args = parser.parse_args()
	train(args)