Hyperparameter optimization with early stopping and gradient clipping #2

[hbandukw]

DATA: NON AUGMENTED / IMBALANCED
METRIC: weighted F1

Best Hyperparameters: {'learning_rate': 0.00011934626028114851, 'dropout_rate': 0.14349283542403013, 'num_trainable_layers': 2, 'train_batch_size': 32, 'grad_clip': 0.5835870650159247, 'step_size': 6, 'gamma': 0.8223727743347635, 'max_length': 319}

I ran it again in my code:
Epoch 1 - Train Loss: 1.2385, Val Loss: 1.0269, Val Accuracy: 0.6440
Epoch 2 - Train Loss: 0.8503, Val Loss: 0.9978, Val Accuracy: 0.6538
Epoch 3 - Train Loss: 0.6682, Val Loss: 1.0414, Val Accuracy: 0.6694
Epoch 4 - Train Loss: 0.4806, Val Loss: 1.1070, Val Accuracy: 0.6645
Epoch 5 - Train Loss: 0.3256, Val Loss: 1.2361, Val Accuracy: 0.6743
Early stopping at epoch 5
Final Model Accuracy: 0.6743

[hbandukw]

import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification, DistilBertConfig
from sklearn.metrics import accuracy_score, f1_score
from sklearn.model_selection import train_test_split
from torch.optim import AdamW, lr_scheduler
import optuna

# Kaggle setup
output_dir = "/kaggle/working/output"
os.makedirs(output_dir, exist_ok=True)

# Load and prepare data
df = pd.read_parquet("/kaggle/input/climatetext/train.parquet")
df['label_int'] = df['label'].str.split("_").str[0].astype('int')
label_dict = df[['label_int', 'label']].drop_duplicates().set_index('label_int')['label'].to_dict()

tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased', do_lower_case=True)

class QuotesDataset(Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx], dtype=torch.long)
        return item

    def __len__(self):
        return len(self.labels)

def encode_data(tokenizer, texts, labels, max_length):
    encodings = tokenizer(texts, truncation=True, padding='max_length', max_length=max_length, return_tensors='pt')
    return QuotesDataset(encodings, labels)

texts = df["quote"].to_list()
labels = df["label_int"].to_list()

X_train, X_val, y_train, y_val = train_test_split(texts, labels, test_size=0.2, random_state=42, stratify=labels)

def objective(trial):
    # Hyperparameters to optimize
    learning_rate = trial.suggest_loguniform('learning_rate', 1e-5, 5e-4)
    dropout_rate = trial.suggest_uniform('dropout_rate', 0.1, 0.5)
    num_trainable_layers = trial.suggest_int('num_trainable_layers', 1, 6)
    train_batch_size = trial.suggest_categorical('train_batch_size', [16, 32, 64])
    grad_clip = trial.suggest_uniform('grad_clip', 0.5, 5.0)
    step_size = trial.suggest_int('step_size', 5, 10)
    gamma = trial.suggest_uniform('gamma', 0.7, 0.9)
    max_length = trial.suggest_int('max_length', 128, 512)

    # Dataset
    train_dataset = encode_data(tokenizer, X_train, y_train, max_length)
    val_dataset = encode_data(tokenizer, X_val, y_val, max_length)

    train_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=train_batch_size, shuffle=False)

    # Model configuration
    config = DistilBertConfig.from_pretrained(
        'distilbert-base-uncased', 
        num_labels=len(label_dict),
        dropout=dropout_rate,
        attention_dropout=dropout_rate
    )
    model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', config=config)

    # Freeze layers except for the last `num_trainable_layers`
    for name, param in model.distilbert.named_parameters():
        param.requires_grad = False
    for layer_idx in range(6 - num_trainable_layers, 6):
        for name, param in model.distilbert.transformer.layer[layer_idx].named_parameters():
            param.requires_grad = True
    for name, param in model.classifier.named_parameters():
        param.requires_grad = True

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

    optimizer = AdamW(model.parameters(), lr=learning_rate)
    scheduler = lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma)

    # Training loop (only 3 epochs for HPO efficiency)
    epochs = 3
    best_val_f1 = 0
    for epoch in range(epochs):
        model.train()
        for batch in train_loader:
            optimizer.zero_grad()
            batch = {k: v.to(device) for k, v in batch.items()}
            outputs = model(**batch)
            loss = outputs.loss
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
            optimizer.step()
        scheduler.step()

        # Validation loop
        model.eval()
        all_predictions, all_true_labels = [], []
        with torch.no_grad():
            for batch in val_loader:
                batch = {k: v.to(device) for k, v in batch.items()}
                outputs = model(**batch)
                predictions = torch.argmax(outputs.logits, dim=-1)
                all_predictions.extend(predictions.cpu().numpy())
                all_true_labels.extend(batch['labels'].cpu().numpy())
        val_f1 = f1_score(all_true_labels, all_predictions, average='weighted')
        best_val_f1 = max(best_val_f1, val_f1)

    return best_val_f1

# Run hyperparameter optimization
study = optuna.create_study(direction='maximize')
study.optimize(objective, n_trials=50)  # Adjust the number of trials as needed

# Save best hyperparameters
print("Best Hyperparameters:", study.best_params)
with open(f"{output_dir}/best_hyperparameters.txt", "w") as f:
    f.write(str(study.best_params))

[hbandukw]

import os
import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification, DistilBertConfig
from sklearn.metrics import f1_score, confusion_matrix, balanced_accuracy_score, precision_score, recall_score, accuracy_score
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
import seaborn as sns
from torch.optim import AdamW, lr_scheduler
import random

# Set seed for reproducibility
torch.manual_seed(42)
np.random.seed(42)
random.seed(42)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

# Training params
MAX_LENGTH = 319
TRAIN_BATCH_SIZE = 32
VAL_BATCH_SIZE = 32
LEARNING_RATE = 0.00011934626028114851
STEP_SIZE = 6
GAMMA = 0.8223727743347635
EPOCHS = 5   
GRAD_CLIP = 0.5835870650159247  
PATIENCE = 3  # early stopping: Number of epochs to wait for improvement
MIN_DELTA = 0.01 # early stopping: Minimum loss improvement required
DROPOUT_RATE = 0.14349283542403013
NUM_TRAINABLE_LAYERS = 2


# Ensure output directory exists
output_dir = "/kaggle/working/output"
os.makedirs(output_dir, exist_ok=True)

# Load and prepare data
df = pd.read_parquet("/kaggle/input/climatetext/train.parquet")
df['label_int'] = df['label'].str.split("_").str[0].astype('int')

# Create dictionary for label names
label_dict = df[['label_int', 'label']].drop_duplicates().set_index('label_int')['label'].to_dict()

# Initialize tokenizer and model
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased', do_lower_case=True)

# Initialize model configuration with custom dropout rate
config = DistilBertConfig.from_pretrained('distilbert-base-uncased', 
                                          num_labels=8, 
                                          dropout=DROPOUT_RATE,
                                          attention_dropout=DROPOUT_RATE)

# Load the model with the updated configuration
model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', config=config)

# Freeze all layers except the last `num_trainable_layers`
for name, param in model.distilbert.named_parameters():
    param.requires_grad = False

# Unfreeze the last `num_trainable_layers`
for layer_idx in range(6 - NUM_TRAINABLE_LAYERS, 6):  # DistilBERT has 6 layers
    for name, param in model.distilbert.transformer.layer[layer_idx].named_parameters():
        param.requires_grad = True

# Ensure classification head is trainable
for name, param in model.classifier.named_parameters():
    param.requires_grad = True

# Dataset and DataLoader preparation
class QuotesDataset(Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx], dtype=torch.long)
        return item

    def __len__(self):
        return len(self.labels)

def encode_data(tokenizer, texts, labels, max_length):
    encodings = tokenizer(texts, truncation=True, padding='max_length', max_length=max_length, return_tensors='pt')
    return QuotesDataset(encodings, labels)

# Initialize early stopping parameters
best_val_loss = float('inf')  # Track the best validation loss
no_improvement_count = 0  # Counter for epochs with no improvement

# Split data
texts = df["quote"].to_list()
labels = df["label_int"].to_list()

X_train, X_val, y_train, y_val = train_test_split(texts, labels, test_size=0.2, random_state=42, stratify=labels)

train_dataset = encode_data(tokenizer, X_train, y_train, MAX_LENGTH)
val_dataset = encode_data(tokenizer, X_val, y_val, MAX_LENGTH)

train_loader = DataLoader(train_dataset, batch_size=TRAIN_BATCH_SIZE, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=VAL_BATCH_SIZE, shuffle=False)

# Initialize optimizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
optimizer = AdamW(model.parameters(), lr=LEARNING_RATE)
scheduler = lr_scheduler.StepLR(optimizer, step_size=STEP_SIZE, gamma=GAMMA)

# Plotting functions
def plot_accuracies(training_accuracies, validation_accuracies):
    plt.figure(figsize=(8, 6))
    plt.plot(training_accuracies, label='Training Accuracy')
    plt.plot(validation_accuracies, label='Validation Accuracy')
    plt.title('Training and Validation Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.legend()
    plt.savefig(f"{output_dir}/accuracy_plot.png")
    plt.close()

def plot_confusion_matrix(cm, class_labels, epoch, output_dir):
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=class_labels, yticklabels=class_labels)
    plt.title(f'Confusion Matrix for Epoch {epoch + 1}')
    plt.ylabel('True Label')
    plt.xlabel('Predicted Label')
    plt.savefig(os.path.join(output_dir, f'confusion_matrix_epoch_{epoch + 1}.png'))
    plt.close()

# Training and validation loop
train_losses = []
val_losses = []
train_accuracies = []
val_accuracies = []
metrics_df = pd.DataFrame()

for epoch in range(EPOCHS):
    model.train()
    train_loss = 0
    correct_train = 0
    total_train = 0

    for batch in train_loader:
        optimizer.zero_grad()
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()

        # Gradient clipping
        # torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)  # Original line
        torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)  # Clip gradients to stabilize training

        optimizer.step()  
        train_loss += loss.item()

        predictions = torch.argmax(outputs.logits, dim=-1)
        correct_train += (predictions == batch['labels']).sum().item()
        total_train += batch['labels'].size(0)

    scheduler.step()
    train_losses.append(train_loss / len(train_loader))
    train_accuracies.append(correct_train / total_train)

    model.eval()
    val_loss = 0
    correct_val = 0
    total_val = 0

    # Clear predictions and true labels before each validation loop
    # all_predictions, all_true_labels = []  # Original line
    all_predictions, all_true_labels = [], []  # Reset for each epoch

    with torch.no_grad():
        for batch in val_loader:
            batch = {k: v.to(device) for k, v in batch.items()}
            outputs = model(**batch)
            val_loss += outputs.loss.item()
            predictions = torch.argmax(outputs.logits, dim=-1)
            all_predictions.extend(predictions.cpu().numpy())
            all_true_labels.extend(batch['labels'].cpu().numpy())

            correct_val += (predictions == batch['labels']).sum().item()
            total_val += batch['labels'].size(0)

    val_losses.append(val_loss / len(val_loader))
    val_accuracies.append(correct_val / total_val)

    # Compute confusion matrix
    cm = confusion_matrix(all_true_labels, all_predictions)
    class_labels = [label_dict.get(i, f'Class {i}') for i in range(len(np.unique(all_true_labels)))]

    # Save confusion matrix only for the best validation loss
    # plot_confusion_matrix(cm, class_labels, epoch, output_dir)  # Original line
    if val_loss == best_val_loss:
        plot_confusion_matrix(cm, class_labels, epoch, output_dir)

    # Print training progress per epoch
    # No print statement included for validation loss and accuracy
    print(f"Epoch {epoch + 1} - Train Loss: {train_losses[-1]:.4f}, Val Loss: {val_losses[-1]:.4f}, Val Accuracy: {val_accuracies[-1]:.4f}")

    # Early stopping check
    if (best_val_loss - val_loss) > MIN_DELTA:  # Significant improvement
        best_val_loss = val_loss
        no_improvement_count = 0

        # Save the best model
        torch.save(model.state_dict(), f"{output_dir}/best_model.pth")
    else:  # No significant improvement
        no_improvement_count += 1
        if no_improvement_count >= PATIENCE:
            print(f"Early stopping at epoch {epoch + 1}")
            break

    # Calculate metrics
    balanced_acc = balanced_accuracy_score(all_true_labels, all_predictions)
    average_f1 = f1_score(all_true_labels, all_predictions, average='macro')
    weighted_f1 = f1_score(all_true_labels, all_predictions, average='weighted')
    f1_scores_per_class = f1_score(all_true_labels, all_predictions, average=None)
    precision_per_class = precision_score(all_true_labels, all_predictions, average=None, zero_division=0)
    recall_per_class = recall_score(all_true_labels, all_predictions, average=None, zero_division=0)

    # Append metrics to DataFrame with class labels
    epoch_metrics = {
        "Epoch": epoch + 1,
        "Train Loss": train_losses[-1],
        "Validation Loss": val_losses[-1],
        "Train Accuracy": train_accuracies[-1],
        "Validation Accuracy": val_accuracies[-1],
        "Balanced Accuracy": balanced_acc,
        "Average F1": average_f1,
        "Weighted F1": weighted_f1,
    }
    epoch_metrics.update({f"{label_dict[i]} F1": f1_scores_per_class[i] for i in range(len(f1_scores_per_class))})
    epoch_metrics.update({f"{label_dict[i]} Precision": precision_per_class[i] for i in range(len(precision_per_class))})
    epoch_metrics.update({f"{label_dict[i]} Recall": recall_per_class[i] for i in range(len(recall_per_class))})

    metrics_df = pd.concat([metrics_df, pd.DataFrame([epoch_metrics])], ignore_index=True)

# Save metrics to CSV and plot accuracies
metrics_df.to_csv(f"{output_dir}/training_metrics.csv", index=False)
plot_accuracies(train_accuracies, val_accuracies)

# Calculate final model accuracy
final_accuracy = accuracy_score(all_true_labels, all_predictions)
print(f"Final Model Accuracy: {final_accuracy:.4f}")

# Save the trained model
model_path = f"{output_dir}/trained_model.pth"
torch.save(model.state_dict(), model_path)
print(f"Model saved to {model_path}")