ANLP_WS24_CA2/BERT.py

214 lines
7.7 KiB
Python

import random
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Subset
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from transformers import BertForSequenceClassification, AutoTokenizer
import numpy as np
from datetime import datetime
import json
import itertools
import Datasets
import dataset_helper
import EarlyStopping
import ml_helper
import ml_history
import ml_train
SEED = 501
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)
torch.backends.cudnn.deterministic = True
class CustomBert(nn.Module):
def __init__(self,dropout):
super().__init__()
#Bert + Custom Layers (Not a tuple any longer -- idk why)
self.bfsc = BertForSequenceClassification.from_pretrained("bert-base-uncased")
self.dropout = nn.Dropout(dropout)
self.classifier = nn.Linear(2,1)
# self.sm = nn.Softmax(dim=1)
def forward(self, input_ids, attention_mask):
x = self.bfsc(input_ids, attention_mask = attention_mask)
x = self.dropout(x[0])
x = self.classifier(x)
x = x.squeeze()
return x
def freeze_bert_params(self):
for param in self.bfsc.named_parameters():
param[1].requires_grad_(False)
def unfreeze_bert_params(self):
for param in self.bfsc.named_parameters():
param[1].requires_grad_(True)
if __name__ == '__main__':
# Hyperparameter und Konfigurationen
params = {
# Training
"epochs": [20],
"patience": [7],
"learning_rate": [1e-5],
"weight_decay": [5e-4],
# Model
"dropout": [0.3]
}
# Generate permutations of hyperparameters
keys, values = zip(*params.items())
grid_params = [dict(zip(keys, v)) for v in itertools.product(*values)]
best_params = {}
best_params_rmse = -1
#Example usage of grid_params
for param_set in grid_params:
print(param_set)
print('Number of grid_params:', len(grid_params))
# Configs
GLOVE_PATH = 'data/glove.6B.100d.txt'
DATA_PATH = 'data/hack.csv'
FREEZE_BERT = False
EMBEDDING_DIM = 100
TEST_SIZE = 0.1
VAL_SIZE = 0.1
MAX_LEN = 128 #280
BATCH_SIZE = 32
N_MODELS = 1
USE_GIRD_SEARCH = False
models = []
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
# Daten laden und vorbereiten
# embedding_matrix, word_index, vocab_size, d_model = dataset_helper.get_embedding_matrix(
# gloVe_path=GLOVE_PATH, emb_len=EMBEDDING_DIM)
X, y = dataset_helper.load_preprocess_data(path_data=DATA_PATH, verbose=True)
# Aufteilen der Daten
data_split = dataset_helper.split_data(X, y, test_size=TEST_SIZE, val_size=VAL_SIZE)
# Initialize BertTokenizer from Pretrained
tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased",do_lower_case=True)
print("Tokenizer Initialized")
# Dataset und DataLoader
train_dataset = Datasets.BertDataset(tokenizer, data_split['train']['X'], data_split['train']['y'], max_len=MAX_LEN)
val_dataset = Datasets.BertDataset(tokenizer, data_split['val']['X'], data_split['val']['y'], max_len=MAX_LEN)
test_dataset = Datasets.BertDataset(tokenizer, data_split['test']['X'], data_split['test']['y'], max_len=MAX_LEN)
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)
subset_size = len(train_dataset) // N_MODELS
device = ml_helper.get_device(verbose=True, include_mps=False)
# assert if N_MODLES > 1, than grid_params should be len 1
if N_MODELS > 1 and len(grid_params) > 1 or N_MODELS > 1 and USE_GIRD_SEARCH:
raise ValueError("If N_MODELS > 1, than grid_params should be len 1")
if not USE_GIRD_SEARCH:
print('Using best params')
# load best params
params_name = f'models/best_params_BERT.json'
with open(params_name, 'r') as f:
best_params = json.load(f)
grid_params = [best_params]
for i in range(N_MODELS):
model_name = f'BERT.pt'
hist_name = f'BERT_history'
if N_MODELS > 1:
model_name = f'BERT_{i}_ensemble.pt'
hist_name = f'BERT_{i}_ensemble_history'
subset_indices = dataset_helper.ensemble_data_idx(train_dataset.labels, N_MODELS, i, methods='bootstrap')
train_dataset_sub = Subset(train_dataset, subset_indices)
train_loader = DataLoader(train_dataset_sub, batch_size=BATCH_SIZE, shuffle=True)
for para_idx, params in enumerate(grid_params):
if len(grid_params) > 1:
model_name = f'BERT_{i}_param_{para_idx}.pt'
hist_name = f'BERT_{i}_param_{para_idx}_history'
model = CustomBert(dropout=params["dropout"])
model = model.to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=params["learning_rate"], weight_decay=params["weight_decay"])
early_stopping = EarlyStopping.EarlyStoppingCallback(patience=params["patience"], verbose=True, model_name=model_name)
hist = ml_history.History()
# Training und Validierung
for epoch in range(params["epochs"]):
ml_train.train_epoch(model, train_loader, criterion, optimizer, device, hist, epoch, params["epochs"], bert_freeze=FREEZE_BERT, is_bert=True)
val_rmse = ml_train.validate_epoch(model, val_loader, epoch, criterion, device, hist, is_bert=True)
early_stopping(val_rmse, model)
if early_stopping.early_stop:
print("Early stopping triggered.")
break
# Load best model
model.load_state_dict(torch.load('models/checkpoints/' + model_name, weights_only=False))
models.append(model)
# Test Evaluation
test_labels, test_preds = ml_train.test_loop(model, test_loader, device, is_bert=True)
hist.add_test_results(test_labels, test_preds)
# save training history
hist.save_history(hist_name, timestamp)
# RMSE, MAE und R²-Score für das Test-Set
test_mae = mean_absolute_error(test_labels, test_preds)
test_rmse = np.sqrt(mean_squared_error(test_labels, test_preds))
test_r2 = r2_score(test_labels, test_preds)
print(f"Test RMSE: {test_rmse:.4f}, Test MAE: {test_mae:.4f}, Test R²: {test_r2:.4f}")
if test_rmse < best_params_rmse:
best_params_rmse = test_rmse
best_params = params
if len(grid_params) > 1:
best_params_name = f'models/best_params_BERT.json'
with open(best_params_name, 'w') as f:
json.dump(best_params, f)
if N_MODELS >1:
# Ensemble Prediction
ensemble_test_preds = ml_train.ensemble_predict(models, test_loader, device, is_bert=True)
ensemble_avg_preds = np.mean(ensemble_test_preds, axis=0)
# Save ensemble predictions as json
ensemble_preds_path = f'histories/ensemble_preds_BERT_{timestamp}.json'
with open(ensemble_preds_path, 'w') as f:
json.dump(ensemble_avg_preds.tolist(), f)
# Test Evaluation
test_labels = test_dataset.labels
test_mse = mean_squared_error(test_labels, ensemble_avg_preds)
test_mae = mean_absolute_error(test_labels, ensemble_avg_preds)
test_r2 = r2_score(test_labels, ensemble_avg_preds)
print(f"Ensemble Test RMSE: {test_mse:.4f}, Test MAE: {test_mae:.4f}, Test R²: {test_r2:.4f}")