continue.py

import os
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from PIL import Image
from transformers import GPT2Tokenizer
import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm
from collections import Counter
from nltk.tokenize import word_tokenize
from sklearn.model_selection import train_test_split
from torchvision import transforms
from model import VQAModel
device = 'cuda'
class Vocab:
    def __init__(self):
        self.vocab = None
        self.vocab_size = None
        self.word2idx = None
        self.idx2word = None
        self.pad = '<pad>'
        self.bos = '<bos>'
        self.eos = '<eos>'
        self.unk = '<unk>'
    def build_vocab(self, df, min_freq=1):
        counter = Counter()
        for ans in df['answer']:
            tokens = word_tokenize(ans.lower())
            counter.update(tokens)
        vocab = sorted([word for word, freq in counter.items() if freq >= min_freq])
        vocab = [self.pad, self.bos, self.eos, self.unk] + vocab
        word2idx = {word: idx for idx, word in enumerate(vocab)}
        idx2word = {idx: word for word, idx in word2idx.items()}
        self.vocab = vocab
        self.word2idx = word2idx
        self.idx2word = idx2word
        self.vocab_size = len(vocab)
        self.pad_token_id = self.word2idx["<pad>"]
        self.bos_token_id = self.word2idx["<bos>"]
        self.eos_token_id = self.word2idx["<eos>"]
        self.unk_token_id = self.word2idx["<unk>"]
    def encoder(self, text, max_len):
        tokens = word_tokenize(text.lower())
        token_ids = [self.word2idx.get(token, self.unk_token_id) for token in tokens]
        token_ids = [self.bos_token_id] + token_ids + [self.eos_token_id]
        if len(token_ids) < max_len:
            token_ids += [self.pad_token_id] * (max_len - len(token_ids))
        else:
            token_ids = token_ids[:max_len]
        return token_ids
    def decoder(self, token_ids):
        tokens = []
        for idx in token_ids:
            if idx == self.eos_token_id:
                break
            if idx in (self.pad_token_id, self.bos_token_id):
                continue
            tokens.append(self.idx2word.get(idx, "<unk>"))
        return ' '.join(tokens).strip()
class AugmentedVQADataset(Dataset):
    def __init__(self, df, img_dir, question_tokenizer, text_processor, clip_processor,
                 question_max_len=32, answer_max_len=16, augment=True):
        self.df = df
        self.img_dir = img_dir
        self.question_tokenizer = question_tokenizer
        self.text_processor = text_processor
        self.clip_processor = clip_processor
        self.question_max_len = question_max_len
        self.answer_max_len = answer_max_len
        self.augment = augment
        if augment:
            self.transform = transforms.Compose([
                transforms.RandomHorizontalFlip(p=0.5),
                transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
                transforms.RandomRotation(10),
            ])
        else:
            self.transform = None
    def __len__(self):
        return len(self.df)
    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        img_path = os.path.join(self.img_dir, row['image_path'])
        image = Image.open(img_path).convert('RGB')
        question = row['question']
        answer = row['answer']
        if self.augment and self.transform:
            image = self.transform(image)
        question_tokenized = self.question_tokenizer(
            question,
            padding='max_length',
            truncation=True,
            max_length=self.question_max_len,
            return_tensors='pt'
        )
        answer_ids = self.text_processor.encoder(answer, max_len=self.answer_max_len)
        image = self.clip_processor(image)
        return {
            'image_path': img_path,
            'image': image,
            'question_ids': question_tokenized['input_ids'].squeeze(0),
            'question_mask': question_tokenized['attention_mask'].squeeze(0),
            'answer_ids': torch.tensor(answer_ids, dtype=torch.long)
        }
def save_checkpoint(model, optimizer, epoch, vocab, path):
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'vocab': vocab.vocab,
        'word2idx': vocab.word2idx,
        'idx2word': vocab.idx2word,
        'pad_token_id': vocab.pad_token_id,
        'bos_token_id': vocab.bos_token_id,
        'eos_token_id': vocab.eos_token_id,
        'unk_token_id': vocab.unk_token_id,
        'question_max_len': model.question_max_len,
        'answer_max_len': model.answer_max_len
    }, path)
def plot_losses(train_losses, val_losses, save_path="loss_plot.png"):
    plt.figure(figsize=(8,6))
    plt.plot(train_losses, label="Train Loss")
    plt.plot(val_losses, label="Validation Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.title("Train vs Validation Loss")
    plt.legend()
    plt.savefig(save_path)
    plt.close()
def train_one_epoch(model, dataloader, optimizer, device, scaler, vocab):
    model.train()
    total_loss = 0
    total_token_acc = 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, :]
            shifted_answers = answers[:, 1:]
            loss = criterion(
                shifted_logits.reshape(-1, shifted_logits.size(-1)),
                shifted_answers.reshape(-1)
            )
            predicted_tokens = shifted_logits.argmax(dim=-1)
            correct = (predicted_tokens == shifted_answers).float()
            mask = (shifted_answers != vocab.pad_token_id).float()
            token_acc = (correct * mask).sum() / mask.sum()
            total_token_acc += token_acc.item()
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optimizer)
        scaler.update()
        total_loss += loss.item()
    avg_loss = total_loss / len(dataloader)
    avg_token_acc = total_token_acc / len(dataloader)
    return avg_loss, avg_token_acc
def validate_one_epoch(model, dataloader, device, vocab):
    model.eval()
    total_loss = 0
    total_token_acc = 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)
            logits = model(images, questions, answer_input_ids=answers)
            shifted_logits = logits[:, :-1, :]
            shifted_answers = answers[:, 1:]
            loss = criterion(
                shifted_logits.reshape(-1, shifted_logits.size(-1)),
                shifted_answers.reshape(-1)
            )
            total_loss += loss.item()
            predicted_tokens = shifted_logits.argmax(dim=-1)
            correct = (predicted_tokens == shifted_answers).float()
            mask = (shifted_answers != vocab.pad_token_id).float()
            token_acc = (correct * mask).sum() / mask.sum()
            total_token_acc += token_acc.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)
    avg_token_acc = total_token_acc / len(dataloader)
    exact_match_acc = exact_matches / total_samples
    return avg_loss, avg_token_acc, exact_match_acc
def main():
    print()
    print("# VQA: Continue Training (Same Settings)")
    print()
    import random
    import numpy as np
    torch.manual_seed(42)
    random.seed(42)
    np.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")
    RESUME_CHECKPOINT = r"./output2/continued_training/vqa_checkpoint.pt"
    OUTPUT_DIR = r"./output2/continued_training_2"
    CHECKPOINT_PATH = 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
    additional_epochs = 50
    patience = 8
    question_max_len = 20
    answer_max_len = 12
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(device)
    print(f"Loading checkpoint from: {RESUME_CHECKPOINT}")
    checkpoint = torch.load(RESUME_CHECKPOINT, map_location=device)
    start_epoch = checkpoint['epoch'] + 1
    metadata = pd.read_csv(CSV_PATH)
    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 Vocab Size: {len(vocab.vocab)}")
    print(f"Resuming from epoch: {start_epoch}")
    train_df, test_df = train_test_split(metadata, 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)}")
    print()
    model = VQAModel(
        vocab_size=len(vocab.vocab),
        device=device,
        question_max_len=question_max_len,
        answer_max_len=answer_max_len,
        pad_token_id=vocab.pad_token_id,
        bos_token_id=vocab.bos_token_id,
        eos_token_id=vocab.eos_token_id,
        unk_token_id=vocab.unk_token_id,
        hidden_size=512,
        num_layers=2
    ).to(device)
    clip_processor = model.clip_preprocess
    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("Model loaded from checkpoint!")
    if model.fine_tuning_mode:
        print("Model already in fine-tuning mode (encoders unfrozen)")
    else:
        print("Continuing with same training configuration")
    print()
    train_dataset = AugmentedVQADataset(
        train_df, DATA_DIR, question_tokenizer, vocab,
        clip_processor=clip_processor,
        question_max_len=question_max_len,
        answer_max_len=answer_max_len,
        augment=True
    )
    val_dataset = AugmentedVQADataset(
        val_df, DATA_DIR, question_tokenizer, vocab,
        clip_processor=clip_processor,
        question_max_len=question_max_len,
        answer_max_len=answer_max_len,
        augment=False
    )
    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)
    trainable_params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.AdamW(trainable_params, lr=1e-6, weight_decay=1e-4)
    print(f"Trainable parameters: {sum(p.numel() for p in trainable_params):,}")
    if 'optimizer_state_dict' in checkpoint:
        try:
            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
            print("Optimizer state loaded from checkpoint!")
            for param_group in optimizer.param_groups:
                print(f"  Loaded LR: {param_group['lr']}")
        except Exception as e:
            print(f"Could not load optimizer state: {e}")
            print("Using fresh optimizer")
    else:
        print("No optimizer state in checkpoint, using fresh optimizer")
    print()
    scaler = torch.amp.GradScaler(device)
    best_val_exact_match = 0.0
    counter = 0
    logs = []
    if os.path.exists(LOG_CSV):
        old_logs = pd.read_csv(LOG_CSV)
        logs = old_logs.values.tolist()
        best_val_exact_match = old_logs['val_exact_match'].max()
        print(f"Previous best exact match: {best_val_exact_match:.4f}")
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='max', factor=0.5, patience=4, verbose=True
    )
    total_epochs = start_epoch + additional_epochs
    for epoch in range(start_epoch, total_epochs):
        print(f"\nEpoch {epoch+1}/{total_epochs}")
        train_loss, train_token_acc = train_one_epoch(model, train_loader, optimizer, device, scaler, vocab)
        val_loss, val_token_acc, val_exact_match = validate_one_epoch(model, val_loader, device, vocab)
        print(f"Train Loss: {train_loss:.4f} | Train Token Acc: {train_token_acc:.4f}")
        print(f"Val Loss: {val_loss:.4f} | Val Token Acc: {val_token_acc:.4f} | Val Exact Match: {val_exact_match:.4f}")
        print(f"LR: {optimizer.param_groups[0]['lr']}")
        scheduler.step(val_exact_match)
        if val_exact_match > best_val_exact_match:
            best_val_exact_match = val_exact_match
            save_checkpoint(model, optimizer, epoch, vocab, CHECKPOINT_PATH)
            print("Checkpoint saved!")
            counter = 0
        else:
            counter += 1
            print(f"No improvement in exact match for {counter} epochs.")
        if counter >= patience:
            print(f"\nEarly stopping after {patience} epochs without improvement")
            break
        logs.append([epoch+1, train_loss, train_token_acc, val_loss, val_token_acc, val_exact_match, optimizer.param_groups[0]['lr']])
    log_df = pd.DataFrame(logs, columns=["epoch","train_loss","train_token_acc","val_loss","val_token_acc","val_exact_match","lr"])
    log_df.to_csv(LOG_CSV, index=False)
    plot_losses([x[1] for x in logs], [x[3] for x in logs], save_path=LOSS_GRAPH_PATH)
    print("\nContinued training complete!")
    print(f"Best exact match accuracy: {best_val_exact_match:.4f}")
if __name__ == "__main__":
    main()