bert.py

# -*- coding: utf-8 -*-
"""BERT.ipynb
"""

########################################
# 1. Imports and Device Setup
########################################
import os
import csv
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.nn.utils import clip_grad_norm_
from tqdm import tqdm
import matplotlib.pyplot as plt

# Transformers and datasets from Hugging Face:
from transformers import AutoTokenizer, AutoModelForSequenceClassification, DataCollatorWithPadding
from datasets import load_dataset

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

########################################
# 2. Load and Prepare Dataset (GLUE/SST-2)
########################################
# Load the SST-2 dataset from the GLUE benchmark.
dataset_name = "imdb"
# dataset = load_dataset("glue", "sst2")
dataset = load_dataset(dataset_name)
print("Loaded dataset: " + dataset_name)

# Use BERT tokenizer.
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

def tokenize_function(example):
    # Tokenize the "sentence" field.
    # return tokenizer(example["sentence"], truncation=True, padding="max_length", max_length=128)
    return tokenizer(example["text"], truncation=True, padding="max_length", max_length=128)

tokenized_train = dataset["train"].map(tokenize_function, batched=True)
#tokenized_val = dataset["validation"].map(tokenize_function, batched=True)
tokenized_val = dataset["test"].map(tokenize_function, batched=True)

# Set the format to PyTorch tensors.
tokenized_train.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
tokenized_val.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])

# Create DataLoaders.
batch_size = 128
train_loader = DataLoader(tokenized_train, batch_size=batch_size, shuffle=True, num_workers=4)
val_loader = DataLoader(tokenized_val, batch_size=batch_size, shuffle=False, num_workers=4)

########################################
# 3. Define Custom Optimizer: FastAdam
########################################
class FastAdam:
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, custom=False, weight_decay=1e-4):
        self.params = list(params)
        self.lr = lr
        self.beta1, self.beta2 = betas
        self.eps = eps
        self.custom = custom
        self.weight_decay = weight_decay
        self.m = [torch.zeros_like(p) for p in self.params]
        self.v = [torch.zeros_like(p) for p in self.params]
        self.prev_grad = [torch.zeros_like(p) for p in self.params]
        self.t = 0

    def zero_grad(self):
        for p in self.params:
            if p.grad is not None:
                p.grad.zero_()

    def step(self):
        self.t += 1
        for p, m, v, prev_g in zip(self.params, self.m, self.v, self.prev_grad):
            if p.grad is None:
                continue
            grad = p.grad
            # Apply decoupled weight decay (AdamW-style)
            if self.weight_decay != 0:
                p.data.mul_(1 - self.lr * self.weight_decay)
            m[:] = self.beta1 * m + (1 - self.beta1) * grad
            v[:] = self.beta2 * v + (1 - self.beta2) * grad.square()
            m_hat = m / (1 - self.beta1 ** self.t)
            v_hat = v / (1 - self.beta2 ** self.t)
            update = -self.lr * m_hat / (v_hat.sqrt() + self.eps)
            if self.custom:
                update = torch.where(update.abs() < prev_g.abs(), update, 0.1 * prev_g)
            p.data.add_(update)
            prev_g[:] = grad

########################################
# 4. Create the Model
########################################
def create_bert():
    # Load a pretrained BERT-base-uncased model for sequence classification (SST-2 has 2 labels).
    model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)
    return model

########################################
# 5. Define the Training Loop (run_experiment)
########################################
def run_experiment(optimizer_name, model_fn, train_loader, test_loader,
                   num_epochs=3, lr=2e-5, momentum=0.9, max_norm=1.0):
    model = model_fn().to(device)
    do_grad_clip = False

    if optimizer_name.lower() == 'fastadam':
        optimizer = FastAdam(model.parameters(), lr=lr, custom=True, weight_decay=1e-4)
    elif optimizer_name.lower() == 'adam':
        optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
    elif optimizer_name.lower() == 'adam_gc':
        optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
        do_grad_clip = True
    else:
        raise ValueError(f"Optimizer '{optimizer_name}' not recognized.")

    history = {"train_loss": [], "train_acc": [], "test_loss": [], "test_acc": []}

    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0
        train_pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs} Training", leave=False)
        for batch in train_pbar:
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            labels = batch["label"].to(device)
            optimizer.zero_grad()
            outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
            loss = outputs.loss
            loss.backward()
            if do_grad_clip:
                clip_grad_norm_(model.parameters(), max_norm)
            optimizer.step()
            running_loss += loss.item() * input_ids.size(0)
            preds = outputs.logits.argmax(dim=1)
            correct += (preds == labels).sum().item()
            total += input_ids.size(0)
            train_pbar.set_postfix(loss=f"{loss.item():.4f}")
        avg_train_loss = running_loss / total
        train_acc = correct / total

        model.eval()
        running_loss_val = 0.0
        correct_val = 0
        total_val = 0
        val_pbar = tqdm(test_loader, desc=f"Epoch {epoch+1}/{num_epochs} Validation", leave=False)
        with torch.no_grad():
            for batch in val_pbar:
                input_ids = batch["input_ids"].to(device)
                attention_mask = batch["attention_mask"].to(device)
                labels = batch["label"].to(device)
                outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
                loss_val = outputs.loss
                running_loss_val += loss_val.item() * input_ids.size(0)
                preds = outputs.logits.argmax(dim=1)
                correct_val += (preds == labels).sum().item()
                total_val += input_ids.size(0)
                val_pbar.set_postfix(loss=f"{loss_val.item():.4f}")
        avg_val_loss = running_loss_val / total_val
        val_acc = correct_val / total_val

        history["train_loss"].append(avg_train_loss)
        history["train_acc"].append(train_acc)
        history["test_loss"].append(avg_val_loss)
        history["test_acc"].append(val_acc)

        print(f"[{optimizer_name.upper()}] Epoch [{epoch+1}/{num_epochs}] - "
              f"Train Loss: {avg_train_loss:.4f}, Train Acc: {train_acc:.4f}, "
              f"Test Loss: {avg_val_loss:.4f}, Test Acc: {val_acc:.4f}")
    return model, history

########################################
# 6. Run Experiments with Multiple Optimizers
########################################
optimizers_list = ['fastadam', 'adam', 'adam_gc']
results = {}
final_models = {}

num_epochs = 25
lr = 2e-5
momentum = 0.9
max_norm = 1.0

for opt_name in optimizers_list:
    print(f"\n*** Training with {opt_name.upper()} ***")
    model, history = run_experiment(
        opt_name,
        create_bert,
        train_loader,
        val_loader,
        num_epochs=num_epochs,
        lr=lr,
        momentum=momentum,
        max_norm=max_norm
    )
    results[opt_name] = history
    final_models[opt_name] = model

# Optionally, print final metrics for each optimizer.
for opt_name in optimizers_list:
    print(f"{opt_name.upper()} final Train Loss: {results[opt_name]['train_loss'][-1]:.4f}, "
          f"Test Loss: {results[opt_name]['test_loss'][-1]:.4f}, "
          f"Train Acc: {results[opt_name]['train_acc'][-1]:.4f}, Test Acc: {results[opt_name]['test_acc'][-1]:.4f}")