AR / train_v4.py

Upload folder using huggingface_hub

f0fc238 verified 5 months ago

7.89 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 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)

	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)