train.py

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms
import wandb
import argparse
import os
from tqdm import tqdm
import numpy as np

from utils.loss import pose_loss

from model.KANLoc import KANLoc
from data.dataset import ImagePoseDataset


def train_one_epoch(model, dataloader, optimizer, device, epoch):
    """
    Train model for one epoch.
    
    Args:
        model: KANLoc model
        dataloader: Training data loader
        optimizer: Optimizer
        device: Target device (cuda/cpu)
        epoch: Current epoch number
        
    Returns:
        avg_loss: Average total loss
        avg_pos_loss: Average position loss
        avg_quat_loss: Average quaternion loss
    """
    model.train()
    total_loss = 0
    total_pos_loss = 0
    total_quat_loss = 0
    
    pbar = tqdm(dataloader, desc=f"Epoch {epoch}")
    for batch_idx, (img1, img2, gt_pose) in enumerate(pbar):
        img1 = img1.to(device)
        img2 = img2.to(device)
        gt_pose = gt_pose.to(device)
        
        # Forward pass
        optimizer.zero_grad()
        pred_pose = model(img1, img2)
        
        # Calculate loss
        loss, pos_loss, quat_loss = pose_loss(pred_pose, gt_pose)
        
        # Backward pass
        loss.backward()
        optimizer.step()
        
        # Accumulate loss
        total_loss += loss.item()
        total_pos_loss += pos_loss.item()
        total_quat_loss += quat_loss.item()
        
        # Update progress bar
        pbar.set_postfix({
            'loss': f'{loss.item():.4f}',
            'pos': f'{pos_loss.item():.4f}',
            'quat': f'{quat_loss.item():.4f}'
        })
        
        # Log to wandb
        if batch_idx % 10 == 0:
            wandb.log({
                'batch_loss': loss.item(),
                'batch_pos_loss': pos_loss.item(),
                'batch_quat_loss': quat_loss.item(),
                'learning_rate': optimizer.param_groups[0]['lr']
            })
    
    avg_loss = total_loss / len(dataloader)
    avg_pos_loss = total_pos_loss / len(dataloader)
    avg_quat_loss = total_quat_loss / len(dataloader)
    
    return avg_loss, avg_pos_loss, avg_quat_loss


def validate(model, dataloader, device):
    """
    Validate model on validation set.
    
    Args:
        model: KANLoc model
        dataloader: Validation data loader
        device: Target device (cuda/cpu)
        
    Returns:
        avg_loss: Average total loss
        avg_pos_loss: Average position loss
        avg_quat_loss: Average quaternion loss
    """
    model.eval()
    total_loss = 0
    total_pos_loss = 0
    total_quat_loss = 0
    
    with torch.no_grad():
        for img1, img2, gt_pose in tqdm(dataloader, desc="Validating"):
            img1 = img1.to(device)
            img2 = img2.to(device)
            gt_pose = gt_pose.to(device)
            
            pred_pose = model(img1, img2)
            loss, pos_loss, quat_loss = pose_loss(pred_pose, gt_pose)
            
            total_loss += loss.item()
            total_pos_loss += pos_loss.item()
            total_quat_loss += quat_loss.item()
    
    avg_loss = total_loss / len(dataloader)
    avg_pos_loss = total_pos_loss / len(dataloader)
    avg_quat_loss = total_quat_loss / len(dataloader)
    
    return avg_loss, avg_pos_loss, avg_quat_loss


def main(args):
    # Initialize wandb
    wandb.init(
        project=args.project_name,
        name=args.run_name,
        config={
            'batch_size': args.batch_size,
            'learning_rate': args.lr,
            'epochs': args.epochs,
            'img_size': args.img_size,
        }
    )
    
    # Set device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    
    # Define image transforms
    transform = transforms.Compose([
        transforms.Resize((args.img_size, args.img_size)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                           std=[0.229, 0.224, 0.225])
    ])
    
    # Create dataset and dataloader
    train_dataset = ImagePoseDataset(
        txt_file=args.train_txt,
        img_dir=args.img_dir,
        transform=transform
    )
    val_dataset = ImagePoseDataset(
        txt_file=args.val_txt,
        img_dir=args.img_dir,
        transform=transform
    )
    
    train_loader = DataLoader(
        train_dataset, 
        batch_size=args.batch_size, 
        shuffle=True, 
        num_workers=args.num_workers,
        pin_memory=True
    )
    val_loader = DataLoader(
        val_dataset,
        batch_size=args.batch_size, 
        shuffle=False, 
        num_workers=args.num_workers,
        pin_memory=True
    )
    
    # Create model
    model = KANLoc(input_dim=256,
                    output_dim=7,
                    fusion_num_heads=8,
                    fusion_method='concat',
                    mechanisms=[None, None, None, "addition"],
                    inv_bottleneck=True).to(device)
    wandb.watch(model, log='all')
    
    # Optimizer and learning rate scheduler
    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
    
    # Create checkpoint directory
    os.makedirs(args.save_dir, exist_ok=True)
    
    best_val_loss = float('inf')
    
    # Training loop
    for epoch in range(1, args.epochs + 1):
        print(f"\n{'='*50}")
        print(f"Epoch {epoch}/{args.epochs}")
        print(f"{'='*50}")
        
        # Train
        train_loss, train_pos_loss, train_quat_loss = train_one_epoch(
            model, train_loader, optimizer, device, epoch
        )
        
        # Validate
        # val_loss, val_pos_loss, val_quat_loss = validate(model, val_loader, device)
        
        # Update learning rate
        scheduler.step()
        
        # Log epoch-level metrics
        wandb.log({
            'epoch': epoch,
            'train_loss': train_loss,
            'train_pos_loss': train_pos_loss,
            'train_quat_loss': train_quat_loss
        })
        
        print(f"\nTrain Loss: {train_loss:.4f} (Pos: {train_pos_loss:.4f}, Quat: {train_quat_loss:.4f})")
        # print(f"Val Loss: {val_loss:.4f} (Pos: {val_pos_loss:.4f}, Quat: {val_quat_loss:.4f})")
        
        # Save checkpoint
        if epoch % args.save_freq == 0:
            checkpoint_path = os.path.join(args.save_dir, f'checkpoint_epoch_{epoch}.pth')
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'scheduler_state_dict': scheduler.state_dict(),
                'train_loss': train_loss
            }, checkpoint_path)
            print(f"Saved checkpoint: {checkpoint_path}")
        
        # Save best model
        if train_loss < best_val_loss:
            best_val_loss = train_loss
            best_model_path = os.path.join(args.save_dir, 'best_model.pth')
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'val_loss': train_loss,
            }, best_model_path)
            print(f"New best model saved! Val Loss: {train_loss:.4f}")
    
    wandb.finish()
    print("\nTraining completed!")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Train pose estimation model')
    
    # Data parameters
    parser.add_argument('--train_txt', type=str, help='Training data txt file path')
    parser.add_argument('--val_txt', type=str, help='Validation data txt file path')
    parser.add_argument('--img_dir', type=str, help='Image directory path')
    parser.add_argument('--img_size', type=int, default=224, help='Image size')
    
    # Training parameters
    parser.add_argument('--batch_size', type=int, default=32, help='Batch size')
    parser.add_argument('--epochs', type=int, default=100, help='Number of training epochs')
    parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate')
    parser.add_argument('--weight_decay', type=float, default=1e-5, help='Weight decay')
    parser.add_argument('--num_workers', type=int, default=4, help='Number of dataloader workers')
    
    # Save parameters
    parser.add_argument('--save_dir', type=str, default='./checkpoints', help='Model save directory')
    parser.add_argument('--save_freq', type=int, default=10, help='Checkpoint save frequency (epochs)')
    
    # WandB parameters
    parser.add_argument('--project_name', type=str, default='kanloc', help='WandB project name')
    parser.add_argument('--run_name', type=str, default=None, help='WandB run name')
    
    args = parser.parse_args()
    
    main(args)