finetune.py

import os
import random
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from transformers import GPT2Tokenizer
from tqdm import tqdm
from sklearn.model_selection import train_test_split
from model import VQAModel
from train import AugmentedVQADataset, Vocab, save_checkpoint, plot_losses
def create_optimizer_with_differential_lr(model, clip_lr=5e-7, gpt_lr=5e-7, other_lr=3e-5):
    clip_params, gpt_params, other_params = [], [], []
    for name, param in model.named_parameters():
        if param.requires_grad:
            if 'clip_model' in name:
                clip_params.append(param)
            elif 'gpt2_model' in name:
                gpt_params.append(param)
            else:
                other_params.append(param)
    optimizer = torch.optim.AdamW([
        {'params': clip_params, 'lr': clip_lr},
        {'params': gpt_params, 'lr': gpt_lr},
        {'params': other_params, 'lr': other_lr}
    ], weight_decay=1e-4)
    print(f"Optimizer: CLIP params: {len(clip_params)}, GPT-2 params: {len(gpt_params)}, Other params: {len(other_params)}")
    return optimizer
def train_one_epoch(model, dataloader, optimizer, device, vocab, scaler):
    model.train()
    total_loss = 0.0
    criterion = nn.CrossEntropyLoss(ignore_index=vocab.pad_token_id, label_smoothing=0.1)
    for batch in tqdm(dataloader):
        optimizer.zero_grad()
        images = batch['image'].to(device)
        questions = {
            'input_ids': batch['question_ids'].to(device),
            'attention_mask': batch['question_mask'].to(device)
        }
        answers = batch['answer_ids'].to(device)
        with torch.amp.autocast(device):
            logits = model(images, questions, answer_input_ids=answers)
            shifted_logits = logits[:, :-1, :].contiguous()
            shifted_answers = answers[:, 1:].contiguous()
            loss = criterion(
                shifted_logits.view(-1, shifted_logits.size(-1)),
                shifted_answers.view(-1)
            )
        if torch.isnan(loss):
            print("NaN loss detected, skipping batch.")
            continue
        scaler.scale(loss).backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optimizer)
        scaler.update()
        total_loss += loss.item()
    return total_loss / len(dataloader)
def validate_one_epoch(model, dataloader, device, vocab):
    model.eval()
    total_loss = 0.0
    exact_matches = 0
    total_samples = 0
    criterion = nn.CrossEntropyLoss(ignore_index=vocab.pad_token_id)
    with torch.no_grad():
        for batch in tqdm(dataloader):
            images = batch['image'].to(device)
            questions = {
                'input_ids': batch['question_ids'].to(device),
                'attention_mask': batch['question_mask'].to(device)
            }
            answers = batch['answer_ids'].to(device)
            with torch.amp.autocast("cuda"):
                logits = model(images, questions, answer_input_ids=answers)
                shifted_logits = logits[:, :-1, :].contiguous()
                shifted_answers = answers[:, 1:].contiguous()
                loss = criterion(
                    shifted_logits.view(-1, shifted_logits.size(-1)),
                    shifted_answers.view(-1)
                )
            total_loss += loss.item()
            generated = model(images, questions)
            for pred, true in zip(generated, answers):
                pred_text = vocab.decoder(pred.cpu().numpy())
                true_text = vocab.decoder(true.cpu().numpy())
                if pred_text.strip() == true_text.strip():
                    exact_matches += 1
                total_samples += 1
    avg_loss = total_loss / len(dataloader)
    exact_match_acc = exact_matches / total_samples
    return avg_loss, exact_match_acc
def filter_spatial_directional_data(df):
    spatial_keywords = [
        'right', 'left', 'above', 'below', 'top', 'bottom',
        'front', 'behind', 'next to', 'beside', 'near',
        'looking', 'facing', 'pointing', 'direction',
        'where is', 'which side', 'what side'
    ]
    directional_answers = [
        'up', 'down', 'left', 'right', 'forward', 'backward',
        'north', 'south', 'east', 'west', 'straight', 'sideways'
    ]
    spatial_mask = df['question'].str.lower().str.contains('|'.join(spatial_keywords), na=False)
    directional_mask = df['answer'].str.lower().str.contains('|'.join(directional_answers), na=False)
    spatial_df = df[spatial_mask | directional_mask].copy()
    print(f"Found {len(spatial_df)} spatial/directional samples out of {len(df)} total")
    return spatial_df
def main():
    print("# VQA: Spatial-Enhanced Fine-Tuning")
    torch.manual_seed(42)
    np.random.seed(42)
    random.seed(42)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(42)
    DATA_DIR = r"./gen_vqa_v2"
    CSV_PATH = os.path.join(DATA_DIR, "metadata.csv")
    PRETRAINED_CHECKPOINT = "./output2/feature_extraction/vqa_checkpoint.pt"
    OUTPUT_DIR = "./output2/spatial_finetuning"
    FINE_TUNED_CHECKPOINT = os.path.join(OUTPUT_DIR, "vqa_checkpoint.pt")
    LOG_CSV = os.path.join(OUTPUT_DIR, "train_log.csv")
    LOSS_GRAPH_PATH = os.path.join(OUTPUT_DIR, "loss_plot.png")
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    batch_size = 64
    num_epochs = 50
    patience = 8
    clip_layers_to_unfreeze = 8
    gpt_layers_to_unfreeze = 8
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    checkpoint = torch.load(PRETRAINED_CHECKPOINT, map_location=device)
    metadata = pd.read_csv(CSV_PATH)
    print(f"\nOriginal dataset size: {len(metadata)}")
    spatial_data = filter_spatial_directional_data(metadata)
    if len(spatial_data) < 1000:
        print(f"\nWARNING: Only {len(spatial_data)} spatial samples found!")
        print("Mixing 70% spatial data with 30% general data for balanced training")
        general_data = metadata[~metadata.index.isin(spatial_data.index)].sample(n=min(len(spatial_data)//2, len(metadata)//3), random_state=42)
        mixed_data = pd.concat([spatial_data, general_data]).sample(frac=1, random_state=42).reset_index(drop=True)
    else:
        print(f"Using {len(spatial_data)} spatial/directional samples")
        mixed_data = spatial_data
    vocab = Vocab()
    vocab.vocab = checkpoint['vocab']
    vocab.vocab_size = len(checkpoint['vocab'])
    vocab.word2idx = checkpoint['word2idx']
    vocab.idx2word = checkpoint['idx2word']
    vocab.pad_token_id = checkpoint['pad_token_id']
    vocab.bos_token_id = checkpoint['bos_token_id']
    vocab.eos_token_id = checkpoint['eos_token_id']
    vocab.unk_token_id = checkpoint['unk_token_id']
    print(f"Answer vocabulary size: {len(vocab.vocab)}")
    model = VQAModel(
        vocab_size=len(checkpoint['vocab']),
        device=device,
        question_max_len=checkpoint.get('question_max_len', 20),
        answer_max_len=checkpoint.get('answer_max_len', 12),
        pad_token_id=checkpoint['pad_token_id'],
        bos_token_id=checkpoint['bos_token_id'],
        eos_token_id=checkpoint['eos_token_id'],
        unk_token_id=checkpoint['unk_token_id'],
        hidden_size=512,
        num_layers=2
    ).to(device)
    question_tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
    if question_tokenizer.pad_token is None:
        question_tokenizer.add_special_tokens({"pad_token": "[PAD]"})
        model.gpt2_model.resize_token_embeddings(len(question_tokenizer))
    model.load_state_dict(checkpoint['model_state_dict'], strict=False)
    print("Pretrained model loaded successfully!\n")
    print(f"UNFREEZING {clip_layers_to_unfreeze} CLIP LAYERS & {gpt_layers_to_unfreeze} GPT-2 LAYERS FOR SPATIAL UNDERSTANDING")
    model.unfreeze_clip_layers(num_layers=clip_layers_to_unfreeze)
    model.unfreeze_gpt2_layers(num_layers=gpt_layers_to_unfreeze)
    train_df, test_df = train_test_split(mixed_data, test_size=0.2, random_state=42)
    val_df, test_df = train_test_split(test_df, test_size=0.5, random_state=42)
    print(f"Train size: {len(train_df)}, Val size: {len(val_df)}, Test size: {len(test_df)}\n")
    train_dataset = AugmentedVQADataset(train_df, DATA_DIR, question_tokenizer, vocab,
                               clip_processor=model.clip_preprocess, augment=True,
                               question_max_len=20, answer_max_len=12)
    val_dataset = AugmentedVQADataset(val_df, DATA_DIR, question_tokenizer, vocab,
                             clip_processor=model.clip_preprocess, augment=False,
                             question_max_len=20, answer_max_len=12)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)
    optimizer = create_optimizer_with_differential_lr(
        model,
        clip_lr=3e-7,
        gpt_lr=3e-7,
        other_lr=2e-5
    )
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=4, verbose=True)
    scaler = torch.amp.GradScaler(device)
    print("\nSTARTING SPATIAL-ENHANCED FINE-TUNING")
    best_val_loss = np.inf
    best_exact_match = 0.0
    logs = []
    counter = 0
    for epoch in range(num_epochs):
        print(f"\nSpatial Fine-tuning Epoch {epoch+1}/{num_epochs}")
        train_loss = train_one_epoch(model, train_loader, optimizer, device, vocab, scaler)
        val_loss, val_exact_match = validate_one_epoch(model, val_loader, device, vocab)
        print(f"Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f} | Val Exact Match: {val_exact_match:.4f} | LR: {optimizer.param_groups[0]['lr']}")
        scheduler.step(val_exact_match)
        if val_exact_match > best_exact_match:
            best_exact_match = val_exact_match
            save_checkpoint(model, optimizer, epoch, vocab, FINE_TUNED_CHECKPOINT)
            print("Checkpoint saved!")
            counter = 0
        else:
            counter += 1
            print(f"No improvement for {counter} epochs.")
        if counter >= patience:
            print(f"\nEarly stopping after {patience} epochs without improvement")
            break
        logs.append([epoch + 1, train_loss, val_loss, val_exact_match, optimizer.param_groups[0]['lr']])
    pd.DataFrame(logs, columns=["epoch", "train_loss", "val_loss", "val_exact_match", "lr"]).to_csv(LOG_CSV, index=False)
    plot_losses([x[1] for x in logs], [x[2] for x in logs], save_path=LOSS_GRAPH_PATH)
    print("\nFINE-TUNING COMPLETE")
    print(f"Best exact match: {best_exact_match:.4f}")
    print(f"Model saved to: {FINE_TUNED_CHECKPOINT}")
if __name__ == "__main__":
    main()