Spaces:
Sleeping
Sleeping
File size: 6,176 Bytes
dd7ff72 039c272 dd7ff72 039c272 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | import os
import yaml
import nltk
import pickle
import torch
import mlflow
import mlflow.pytorch
from tqdm import tqdm
from preprocess_data import VQADataset
from nltk.tokenize import word_tokenize
from torch.utils.data import DataLoader
from model import load_model_processor
from transformers import get_linear_schedule_with_warmup
from nltk.translate.bleu_score import sentence_bleu , SmoothingFunction
from mlflow.models import infer_signature
nltk.download('all')
config = yaml.safe_load(open("./config.yaml", "r"))["data_location"]
model_config = yaml.safe_load(open("./config.yaml", "r"))["finetune_model"]
params = yaml.safe_load(open("./param.yaml", "r"))["params"]
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model, processor = load_model_processor()
print(f"Loading processed data from {config['train_processed_data']} and {config['test_processed_data']}.")
with open(f'./{config["train_processed_data"]}', "rb") as f:
train_dataset = pickle.load(f)
with open(f'./{config["test_processed_data"]}', "rb") as f:
test_dataset = pickle.load(f)
print(f"Loaded processed data successfully!!!")
print(f"Length of train dataset is {len(train_dataset)} and test dataset is {len(test_dataset)}")
batch_size = params["batch_size"]
num_epochs = params["num_epochs"]
patience = params["patience"]
gradient_accumulation_steps = params["gradient_accumulation_steps"]
print("Preparing dataloaders................")
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, pin_memory=True)
valid_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, pin_memory=True)
print("Dataloaders prepared successfully!!!")
mlflow.set_experiment("VQA_Model_Training")
min_bleu_score = 0
early_stopping_hook = 0
tracking_information = []
optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=float(params["learning_rate"]), weight_decay=float(params["weight_decay"]))
total_steps = len(train_dataloader) * num_epochs
warmup_steps = total_steps // 10
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps)
scaler = torch.amp.GradScaler('cuda')
with mlflow.start_run():
for epoch in range(num_epochs):
model.train()
epoch_loss = 0
for step, batch in enumerate(tqdm(train_dataloader, desc=f"Epoch {epoch+1} Training")):
input_ids = batch.pop('input_ids').to(device)
pixel_values = batch.pop('pixel_values').to(device)
attention_mask = batch.pop('attention_mask').to(device)
labels = batch.pop('labels').to(device)
with torch.amp.autocast('cuda', dtype=torch.float16):
outputs = model(input_ids=input_ids, pixel_values=pixel_values, attention_mask=attention_mask, labels=labels)
loss = outputs.loss / gradient_accumulation_steps
scaler.scale(loss).backward()
if (step + 1) % gradient_accumulation_steps == 0 or (step + 1) == len(train_dataloader):
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
epoch_loss += loss.item() * gradient_accumulation_steps
# Validation Loop
model.eval()
eval_loss = 0
bleu_scores = []
smooth_fn = SmoothingFunction().method1
with torch.no_grad():
for batch in tqdm(valid_dataloader, desc=f"Epoch {epoch+1} Validating"):
input_ids = batch.pop('input_ids').to(device)
pixel_values = batch.pop('pixel_values').to(device)
attention_mask = batch.pop('attention_mask').to(device)
labels = batch.pop('labels').to(device)
with torch.amp.autocast('cuda', dtype=torch.float16):
outputs = model(input_ids=input_ids, pixel_values=pixel_values, attention_mask=attention_mask, labels=labels)
eval_loss += outputs.loss.item()
generated_ids = model.generate(input_ids=input_ids, pixel_values=pixel_values, attention_mask=attention_mask, max_length=8)
predictions = processor.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
references = processor.tokenizer.batch_decode(labels, skip_special_tokens=True)
for pred, ref in zip(predictions, references):
bleu_scores.append(sentence_bleu([word_tokenize(ref)], word_tokenize(pred), smoothing_function=smooth_fn))
avg_train_loss = epoch_loss / len(train_dataloader)
avg_eval_loss = eval_loss / len(valid_dataloader)
avg_bleu_score = sum(bleu_scores) / len(bleu_scores)
tracking_information.append((avg_train_loss, avg_eval_loss, avg_bleu_score, optimizer.param_groups[0]["lr"]))
mlflow.log_metrics({"train_loss": avg_train_loss, "eval_loss": avg_eval_loss, "bleu_score": avg_bleu_score}, step=epoch)
mlflow.log_param(f"learning_rate_{epoch}", optimizer.param_groups[0]['lr'])
print(f"Epoch {epoch+1} - Train Loss: {avg_train_loss:.4f} - Eval Loss: {avg_eval_loss:.4f} - BLEU Score: {avg_bleu_score:.4f} - LR: {optimizer.param_groups[0]['lr']}")
if avg_bleu_score >= min_bleu_score:
model.save_pretrained(model_config["best"], from_pt=True)
print(f"Model improved (BLEU: {avg_bleu_score:.4f})! Saved to {model_config['best']}")
min_bleu_score = avg_bleu_score
early_stopping_hook = 0
else:
early_stopping_hook += 1
if early_stopping_hook > patience:
print("Early stopping triggered.")
break
model.save_pretrained(model_config["last"], from_pt=True)
scheduler.step()
# Log the saved PEFT adapters as an artifact instead of trying to pickle the Quantized model
mlflow.log_artifacts(model_config["best"], artifact_path="VQA_model_best")
print("Model adapters logged with MLflow.")
|