Anul 3 Semestrul 1
This commit is contained in:
@@ -0,0 +1,220 @@
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
import argparse
|
||||
import logging
|
||||
import json
|
||||
from tqdm import tqdm
|
||||
|
||||
import torch
|
||||
from torch.utils import data
|
||||
from sklearn.metrics import f1_score, precision_score, recall_score, confusion_matrix
|
||||
import numpy as np
|
||||
|
||||
from data_loader import DFGenerator, StoryPointDataset
|
||||
from models import get_model_and_tokenizer
|
||||
from utils import serialize_metrics
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
|
||||
parser = argparse.ArgumentParser(description='Model training details')
|
||||
|
||||
# Dataset details
|
||||
parser.add_argument('--eval_split', type=float, default=0.1, help='Evaluation split')
|
||||
parser.add_argument('--random_seed', type=int, default=42, help='Random seed')
|
||||
parser.add_argument('--max_length', type=int, default=128, help='Maximum number of tokens in input sequence')
|
||||
|
||||
# Model training details
|
||||
parser.add_argument('--checkpoint', type=str, default=None, help='Checkpoint to resume training')
|
||||
parser.add_argument('--epochs', type=int, default=20, help='Number of epochs')
|
||||
parser.add_argument('--batch_size', type=int, default=16, help='Batch size')
|
||||
parser.add_argument('--learning_rate', type=float, default=1e-5, help='Learning rate')
|
||||
parser.add_argument('--weight_decay', type=float, default=1e-5, help='Weight decay')
|
||||
parser.add_argument('--eval', type=bool, default=True, help='Evaluate model after each training epoch')
|
||||
|
||||
# Model saving details
|
||||
parser.add_argument('--experiment_name', type=str, default='storypoint_estimator', help='Experiment name (for checkpoint naming)')
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
class Trainer():
|
||||
|
||||
def __init__(self, args):
|
||||
|
||||
self.experiment_dir = f'experiments/{args.experiment_name}'
|
||||
self.checkpoints_dir = f'{self.experiment_dir}/checkpoints'
|
||||
|
||||
os.makedirs(self.experiment_dir, exist_ok=True)
|
||||
os.makedirs(self.checkpoints_dir, exist_ok=True)
|
||||
|
||||
log_path = f'{self.experiment_dir}/{args.experiment_name}.log'
|
||||
|
||||
self.logger = logging.getLogger('training_pipeline')
|
||||
self.logger.setLevel(logging.INFO)
|
||||
|
||||
file_handler = logging.FileHandler(log_path, mode='a')
|
||||
file_handler.setLevel(logging.INFO)
|
||||
|
||||
self.logger.addHandler(file_handler)
|
||||
|
||||
params_path = f'{self.experiment_dir}/params.json'
|
||||
|
||||
with open(params_path, 'w') as f:
|
||||
json.dump(vars(args), f, indent=4)
|
||||
|
||||
|
||||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
self.model, self.tokenizer = get_model_and_tokenizer(args.checkpoint)
|
||||
self.model.to(self.device)
|
||||
|
||||
df_generator = DFGenerator(eval_split=args.eval_split, random_seed=args.random_seed)
|
||||
|
||||
train_df, val_df = df_generator.create_dataframes()
|
||||
if train_df is None:
|
||||
raise ValueError('No training data available')
|
||||
|
||||
train_dataset = StoryPointDataset(train_df, self.tokenizer, max_length=args.max_length)
|
||||
val_dataset = StoryPointDataset(val_df, self.tokenizer, max_length=args.max_length) if val_df is not None else None
|
||||
|
||||
self.train_loader = data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
|
||||
self.val_loader = data.DataLoader(val_dataset, batch_size=1) if val_dataset is not None else None
|
||||
|
||||
|
||||
self.epochs = args.epochs
|
||||
self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)
|
||||
self.lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=self.epochs, eta_min=1e-7)
|
||||
|
||||
|
||||
self.eval = args.eval
|
||||
self.experiment_name = args.experiment_name
|
||||
|
||||
|
||||
def train(self):
|
||||
|
||||
for epoch in range(self.epochs):
|
||||
self.model.train()
|
||||
correct_predictions = 0
|
||||
total_predictions = 0
|
||||
total_loss = 0
|
||||
|
||||
all_labels = []
|
||||
all_predictions = []
|
||||
|
||||
for batch in tqdm(self.train_loader, desc=f'Epoch {epoch+1}/{self.epochs}'):
|
||||
inputs = {k: v.to(self.model.device) for k, v in batch.items()}
|
||||
outputs = self.model(**inputs)
|
||||
|
||||
loss = outputs.loss
|
||||
loss.backward()
|
||||
|
||||
self.optimizer.step()
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
total_loss += loss.item()
|
||||
_, predicted = torch.max(outputs.logits, 1)
|
||||
|
||||
correct_predictions += (predicted == batch['labels'].to(self.model.device)).sum().item()
|
||||
total_predictions += len(batch['labels'])
|
||||
|
||||
all_labels.extend(batch['labels'].cpu().numpy())
|
||||
all_predictions.extend(predicted.cpu().numpy())
|
||||
|
||||
all_labels = np.array(all_labels)
|
||||
all_predictions = np.array(all_predictions)
|
||||
|
||||
epoch_lr = self.optimizer.param_groups[0]['lr']
|
||||
self.lr_scheduler.step()
|
||||
|
||||
average_loss = total_loss / len(self.train_loader)
|
||||
accuracy = correct_predictions / total_predictions
|
||||
macro_f1 = f1_score(all_labels, all_predictions, average='macro')
|
||||
weighted_f1 = f1_score(all_labels, all_predictions, average='weighted')
|
||||
precision = precision_score(all_labels, all_predictions, average='macro')
|
||||
recall = recall_score(all_labels, all_predictions, average='macro')
|
||||
conf_matrix = confusion_matrix(all_labels, all_predictions)
|
||||
|
||||
metrics = (accuracy, macro_f1, weighted_f1, precision, recall, conf_matrix)
|
||||
|
||||
|
||||
epoch_message = f"Training - Epoch {epoch+1}/{self.epochs}: Loss - {average_loss:.4f}\t\tLR - {epoch_lr:.8f}\n"
|
||||
|
||||
print(epoch_message)
|
||||
self.logger.info(epoch_message)
|
||||
|
||||
|
||||
if self.eval and self.val_loader:
|
||||
self.evaluate(epoch, metrics)
|
||||
|
||||
else:
|
||||
self.save_checkpoint('train', metrics, epoch)
|
||||
|
||||
|
||||
|
||||
def evaluate(self, epoch, train_metrics = None):
|
||||
|
||||
if self.val_loader is None:
|
||||
return
|
||||
|
||||
self.model.eval()
|
||||
correct_predictions = 0
|
||||
total_predictions = 0
|
||||
|
||||
all_labels = []
|
||||
all_predictions = []
|
||||
|
||||
for batch in tqdm(self.val_loader, desc=f'Validation - Epoch {epoch+1}/{self.epochs}'):
|
||||
inputs = {k: v.to(self.model.device) for k, v in batch.items()}
|
||||
with torch.no_grad():
|
||||
outputs = self.model(**inputs)
|
||||
|
||||
_, predicted = torch.max(outputs.logits, 1)
|
||||
correct_predictions += (predicted == batch['labels'].to(self.model.device)).sum().item()
|
||||
total_predictions += len(batch['labels'])
|
||||
|
||||
all_labels.extend(batch['labels'].cpu().numpy())
|
||||
all_predictions.extend(predicted.cpu().numpy())
|
||||
|
||||
all_labels = np.array(all_labels)
|
||||
all_predictions = np.array(all_predictions)
|
||||
|
||||
accuracy = correct_predictions / total_predictions
|
||||
macro_f1 = f1_score(all_labels, all_predictions, average='macro')
|
||||
weighted_f1 = f1_score(all_labels, all_predictions, average='weighted')
|
||||
precision = precision_score(all_labels, all_predictions, average='macro')
|
||||
recall = recall_score(all_labels, all_predictions, average='macro')
|
||||
conf_matrix = confusion_matrix(all_labels, all_predictions)
|
||||
|
||||
metrics = (accuracy, macro_f1, weighted_f1, precision, recall, conf_matrix)
|
||||
|
||||
if train_metrics is not None:
|
||||
all_metrics = list(train_metrics)
|
||||
all_metrics.extend(metrics)
|
||||
metrics = tuple(all_metrics)
|
||||
|
||||
self.save_checkpoint('both', metrics, epoch)
|
||||
|
||||
else:
|
||||
self.save_checkpoint('val', metrics, epoch)
|
||||
|
||||
|
||||
def save_checkpoint(self, mode, metrics, epoch=0):
|
||||
|
||||
json_data = serialize_metrics(mode, metrics)
|
||||
|
||||
checkpoint_dir = f'{self.checkpoints_dir}/epoch_{epoch+1}'
|
||||
os.makedirs(checkpoint_dir, exist_ok=True)
|
||||
|
||||
torch.save(self.model.state_dict(), f'{checkpoint_dir}/{self.experiment_name}_{epoch+1}.pth')
|
||||
with open(f'{checkpoint_dir}/metrics.json', 'w') as f:
|
||||
f.write(json_data)
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
load_dotenv()
|
||||
args = parse_args()
|
||||
|
||||
trainer = Trainer(args)
|
||||
trainer.train()
|
||||
Reference in New Issue
Block a user