import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import Dataset, DataLoader, random_split, Subset from torchvision import transforms import os import numpy as np import h5py from tqdm import tqdm from models import CNNModel_Small, CNNModel_Medium, CNNModel_Large DATA_FILE = "data/book_dataset.h5" MODEL_OUTPUT_DIR = "models/saved_weights_finetuned/" EPOCHS = 3 BATCH_SIZE = 256 LEARNING_RATE = 0.001 VALIDATION_SPLIT = 0.15 EARLY_STOPPING_PATIENCE = 2 DEVICE = torch.device("mps" if torch.backends.mps.is_available() else "cpu") class EarlyStopping: def __init__(self, patience=2, path='checkpoint.pth'): self.patience, self.path = patience, path self.counter, self.val_loss_min = 0, np.inf self.best_score, self.early_stop = None, False def __call__(self, val_loss, model): score = -val_loss if self.best_score is None: self.best_score = score self.save_checkpoint(val_loss, model) elif score < self.best_score: self.counter += 1 if self.counter >= self.patience: self.early_stop = True else: self.best_score = score self.save_checkpoint(val_loss, model) self.counter = 0 def save_checkpoint(self, val_loss, model): print(f'Val loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}). Saving model...') torch.save(model.state_dict(), self.path) self.val_loss_min = val_loss class LabelRemapper: def __init__(self, remap_dict): self.remap_dict = remap_dict def __call__(self, label): return self.remap_dict.get(label, -1) class HDF5Dataset(Dataset): def __init__(self, h5_path, transform=None, target_transform=None): self.h5_path = h5_path self.transform = transform self.target_transform = target_transform with h5py.File(self.h5_path, 'r') as hf: self.length = len(hf['labels']) def __len__(self): return self.length def __getitem__(self, idx): with h5py.File(self.h5_path, 'r') as hf: image = hf['images'][idx] label = hf['labels'][idx] if self.transform: image = self.transform(image) if self.target_transform: label = self.target_transform(label) return image, label def prepare_dataloaders(): print("Preparing dataloaders...") image_transform = transforms.Compose([transforms.ToTensor()]) with h5py.File(DATA_FILE, 'r') as hf: all_labels = hf['labels'][:] unique_labels = np.unique(all_labels) triage_remap = {code: (0 if 48 <= code <= 57 else 1 if 65 <= code <= 90 else 2) for code in unique_labels} triage_target_transform = LabelRemapper(triage_remap) triage_dataset = HDF5Dataset(DATA_FILE, transform=image_transform, target_transform=triage_target_transform) val_size = int(len(triage_dataset) * VALIDATION_SPLIT) train_size = len(triage_dataset) - val_size triage_train, triage_val = random_split(triage_dataset, [train_size, val_size]) datasets = {'triage': (triage_train, triage_val)} expert_filters = { 'digits': (lambda c: 48 <= c <= 57), 'uppercase': (lambda c: 65 <= c <= 90), 'lowercase': (lambda c: 97 <= c <= 122) } for name, condition in expert_filters.items(): class_codes = sorted([c for c in unique_labels if condition(c)]) expert_remap = {code: i for i, code in enumerate(class_codes)} expert_target_transform = LabelRemapper(expert_remap) indices = [i for i, code in enumerate(all_labels) if condition(code)] expert_full_dataset = HDF5Dataset(DATA_FILE, transform=image_transform, target_transform=expert_target_transform) expert_subset = Subset(expert_full_dataset, indices) val_expert_size = int(len(expert_subset) * VALIDATION_SPLIT) train_expert_size = len(expert_subset) - val_expert_size expert_train, expert_val = random_split(expert_subset, [train_expert_size, val_expert_size]) datasets[name] = (expert_train, expert_val) print("Dataloaders are ready.") return datasets def train_model(model_name, model, train_dataset, val_dataset): print(f"\n{'=' * 60}\nTraining: {model_name.upper()} MODEL on {len(train_dataset)} samples\n{'=' * 60}") train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2) val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2) optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE) criterion = nn.CrossEntropyLoss() model_save_path = os.path.join(MODEL_OUTPUT_DIR, f"{model_name}_model_finetuned.pth") early_stopper = EarlyStopping(patience=EARLY_STOPPING_PATIENCE, path=model_save_path) for epoch in range(EPOCHS): model.train() train_loss = 0.0 for images, labels in tqdm(train_loader, desc=f"Epoch {epoch + 1}/{EPOCHS} [Train]"): images, labels = images.to(DEVICE), labels.to(DEVICE, non_blocking=True) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() train_loss += loss.item() * images.size(0) model.eval() val_loss = 0.0 with torch.no_grad(): for images, labels in tqdm(val_loader, desc=f"Epoch {epoch + 1}/{EPOCHS} [Val]"): images, labels = images.to(DEVICE), labels.to(DEVICE, non_blocking=True) outputs = model(images) loss = criterion(outputs, labels) val_loss += loss.item() * images.size(0) avg_train_loss = train_loss / len(train_dataset) avg_val_loss = val_loss / len(val_dataset) print(f"Epoch {epoch + 1}/{EPOCHS} | Train Loss: {avg_train_loss:.6f} | Validation Loss: {avg_val_loss:.6f}") early_stopper(avg_val_loss, model) if early_stopper.early_stop: print("Early stopping triggered.") break print(f"Training finished. Model saved to {model_save_path}") def main(): if not os.path.exists(MODEL_OUTPUT_DIR): os.makedirs(MODEL_OUTPUT_DIR) print(f"Using device: {DEVICE}") datasets = prepare_dataloaders() models_to_train = { 'triage': CNNModel_Large(num_classes=3), 'digits': CNNModel_Small(num_classes=10), 'uppercase': CNNModel_Medium(num_classes=26), 'lowercase': CNNModel_Medium(num_classes=26) } for name, model in models_to_train.items(): model_path = os.path.join(MODEL_OUTPUT_DIR, f"{name}_model_finetuned.pth") if os.path.exists(model_path): print(f"Skipping training for '{name}': Model already exists at {model_path}") continue train_d, val_d = datasets[name] model.to(DEVICE) train_model(name, model, train_d, val_d) if __name__ == "__main__": main()