glisicstefan/age-estimation-resnet50

1import matplotlib.pyplot as plt2import torch3from pathlib import Path4import random5 6def plot_random_predictions(model, dataloader, device, n=4):7    """Choosing random examples from dataloader and shows predictions"""8 9    images, labels = next(iter(dataloader))10    images, labels = images.to(device), labels.to(device)11 12    model.eval()13    with torch.inference_mode():14        preds = model(images).squeeze()15 16    batch_size = images.shape[0]17    indices = random.sample(range(batch_size), k=min(n, batch_size))18 19    plt.figure(figsize=(15, 5))20    for i, idx in enumerate(indices):21        plt.subplot(1, n, i+1)22        23        img = images[idx].cpu().permute(1, 2, 0).numpy()24        img = img * [0.229, 0.224, 0.225] + [0.485, 0.456, 0.406]25        26        plt.imshow(img.clip(0, 1))27        plt.title(f"Real: {labels[idx].item():.0f}y\nPredicted: {preds[idx].item():.1f}y")28        plt.axis("off")29    plt.show()30 31 32def save_model(model, model_name):33 34    target_dir_path = Path("../models")35    target_dir_path.mkdir(parents=True, exist_ok=True) 36    37    model_save_path = target_dir_path / model_name38 39    print(f"[INFO] Saving model to: {model_save_path}")40    torch.save(obj=model.state_dict(), f=model_save_path)41 42 43def plot_loss_curves(results):44    """Shows loss curves (loss) and metric (MAE) for train and test set."""45    46    loss = results["train_loss"]47    test_loss = results["test_loss"]48 49    mae = results["train_mae"]50    test_mae = results["test_mae"]51 52    epochs = range(len(results["train_loss"]))53 54    plt.figure(figsize=(15, 5))55 56    # Grafikon 1: Loss57    plt.subplot(1, 2, 1)58    plt.plot(epochs, loss, label="Train Loss")59    plt.plot(epochs, test_loss, label="Test Loss")60    plt.title("Loss (L1 / MAE)")61    plt.xlabel("Epochs")62    plt.legend()63 64    # Grafikon 2: MAE65    plt.subplot(1, 2, 2)66    plt.plot(epochs, mae, label="Train MAE")67    plt.plot(epochs, test_mae, label="Test MAE")68    plt.title("Mean Absolute Error")69    plt.xlabel("Epochs")70    plt.legend()71    72    plt.show()73