

# # Finetune a ResNet image classifier
# ## Load data
!curl -L https://storage.googleapis.com/aiolympiadmy/ioai-2025-tsp/finetuning-resnet.zip -o data.zip
import zipfile

with zipfile.ZipFile('data.zip', 'r') as zip_ref:
    zip_ref.extractall()
# ## Establishing a baseline
# Load an ImageNet pre-trained ResNet34, and check it's performance on the images in `data/test`.  
# 
# Use accuracy, precision and recall as your metrics for performance.
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np
import random
from torchvision import datasets, transforms, models
from torchvision.models import ResNet34_Weights
from torch.utils.data import DataLoader, Subset, ConcatDataset, Sampler
from tqdm import tqdm
from sklearn.metrics import accuracy_score, precision_score, recall_score
# ### Utilities
import os
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
from collections import Counter
 
def count_images(data_dir):
    counts = {}
    for split in ['train', 'test']:
        split_path = os.path.join(data_dir, split)
        counts[split] = {}
       
        for cls in os.listdir(split_path):
            cls_path = os.path.join(split_path, cls)
            counts[split][cls] = len(os.listdir(cls_path))
   
    return counts
 
data_dir = "./data"
counts = count_images(data_dir)
 
print(counts)
def plot_class_distribution(counts):
    for split in counts:
        classes = list(counts[split].keys())
        values = list(counts[split].values())

        print(classes, values)

        plt.figure()
        plt.bar(classes, values)
        plt.title(f"{split.capitalize()} Class Distribution")
        plt.xlabel("Class")
        plt.ylabel("Number of Images")
        plt.show()

plot_class_distribution(counts)

import random

def show_samples(data_dir, split='train', n=5):
    plt.figure(figsize=(10, 5))
    
    i = 1
    for cls in ['chihuahua', 'muffin']:
        cls_path = os.path.join(data_dir, split, cls)
        images = random.sample(os.listdir(cls_path), n)
        
        for img_name in images:
            img_path = os.path.join(cls_path, img_name)
            img = Image.open(img_path)
            
            plt.subplot(2, n, i)
            plt.imshow(img)
            plt.title(cls)
            plt.axis('off')
            i += 1
    
    plt.tight_layout()
    plt.show()

show_samples(data_dir, "train", 7)
class FixedOrderSampler(Sampler):
    __slots = ("indices", )
    
    def __init__(self, indices):
        self.indices = indices

    def __iter__(self):
        return iter(self.indices)

    def __len__(self):
        return len(self.indices)
from torchvision.transforms import AugMix
from sklearn.model_selection import StratifiedKFold
from torch.utils.data import Subset
from torchvision.transforms import v2
class BinaryImageClassifier:
    __slots__ = ("__lr", "__l2_penalty", "pretrained", "device", "criterion", "train_loss_history", "test_loss_history",
                 "train_metrics_history", "test_results", "model", "optimizer", "train_loader", "test_loader", "train_dir")

    def __init__(self, *, lr, l2_penalty=0, pretrained=True, train_dir="./data/train", test_dir="./data/test"):
        self.__lr = lr
        self.__l2_penalty = l2_penalty
        self.pretrained = pretrained
        self.train_dir = train_dir
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.criterion = nn.BCELoss()  # Binary Cross Entropy Loss
        self.train_loss_history = []
        self.test_loss_history = []
        self.train_metrics_history = []
        self.test_results = {}

        self.reload_model()

        self.train_loader = self.load_data(train_dir, is_train=True)
        self.test_loader = self.load_data(test_dir, is_train=False)


    @property
    def lr(self):
        return self.__lr

    @lr.setter
    def lr(self, val):
        self.__lr = val
        self.__reload_optimizer()


    @property
    def l2_penalty(self):
        return self.__l2_penalty

    @l2_penalty.setter
    def l2_penalty(self, val):
        self.__l2_penalty = val
        self.__reload_optimizer()


    def __reload_optimizer(self):
        self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr, weight_decay=self.l2_penalty)

    def fit_kfold(self, n_epochs, n_splits=5):
        dataset = self.train_loader.dataset  
        X = np.arange(len(dataset))
        y = np.array(dataset.targets)
        skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)
        results = []
        for fold, (train_idx, val_idx) in enumerate(skf.split(X, y), 1):
            print(f"Fold {fold}")
            self.train_loader = DataLoader(Subset(dataset, train_idx),batch_size=64,shuffle=True)
            self.test_loader = DataLoader(Subset(dataset, val_idx),batch_size=64,shuffle=False)
            self.fit(n_epochs=n_epochs, reset=True)
            res = self.predict(show_progress=False)
            results.append(res["metrics"])
            print(res["metrics"])

    def load_data(self, dir_, start=0, stop=None, step=1, is_train=True):
        self.__set_seed()

        if is_train:
            transform = transforms.Compose([
                transforms.Resize((224, 224)),
                transforms.Lambda(lambda img: img.convert("RGB")),
                AugMix(),
                transforms.ToTensor()
            ])
        else:
            transform = transforms.Compose([
                transforms.Resize((224, 224)),
                transforms.Lambda(lambda img: img.convert("RGB")),
                transforms.ToTensor()
            ])
        dataset = datasets.ImageFolder(root=dir_, transform=transform)

        # Shuffle the dataset
        indices = list(range(len(dataset)))
        random.shuffle(indices)
        indices = indices[start:stop:step]

        loader = DataLoader(dataset=dataset, batch_size=64, sampler=FixedOrderSampler(indices))
        return loader


    @staticmethod
    def __set_seed():
        seed = 42

        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed(seed)
            torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


    def reload_model(self):
        self.__set_seed()

        self.model = models.resnet34(weights=ResNet34_Weights.DEFAULT if self.pretrained else None).to(self.device)
        self.model.fc = nn.Sequential(
            nn.Linear(self.model.fc.in_features, 1),  # One final output
            nn.Sigmoid()
        ).to(self.device)
        self.__reload_optimizer()


    @staticmethod
    def __measure_metrics(y_true, y_pred):
        res = {}
        for metric_name, metric_fn in [("accuracy", accuracy_score),
                                       ("precision", lambda *args, **kwargs: precision_score(*args, **kwargs, zero_division=0)),  # to suppress warning
                                       ("recall", lambda *args, **kwargs: recall_score(*args, **kwargs, zero_division=0))]:
            res[metric_name] = metric_fn(y_true, y_pred)
        return res


    def fit(self, n_epochs, *, reset=True, show_progress=True):
        if reset:
            self.train_loss_history = []
            self.test_loss_history = []
            self.train_metrics_history = []
            self.reload_model()
        n_batches = len(self.train_loader)

        self.__set_seed()
        #self.model.train()  # Set model to training mode
        with tqdm(range(n_epochs * n_batches), desc="Training", ncols=150, disable=not show_progress) as pbar:
            for epoch in range(1, n_epochs + 1):
                for batch_idx, (batch_images, batch_labels) in enumerate(self.train_loader, 1):
                    batch_images, batch_labels = batch_images.to(self.device), batch_labels.float().reshape(-1, 1).to(self.device)

                    # Forward pass
                    self.model.train()  # Set model to training mode
                    batch_probas = self.model(batch_images)
                    batch_preds = batch_probas >= 0.5
                    batch_loss = self.criterion(batch_probas, batch_labels)

                    # Backward pass
                    self.optimizer.zero_grad()
                    batch_loss.backward()
                    self.optimizer.step()

                    # Record
                    self.train_metrics_history.append(self.__measure_metrics(batch_labels.cpu(), batch_preds.cpu()))
                    self.train_loss_history.append(batch_loss.item())
                    pbar.set_postfix(loss=f"{self.train_loss_history[-1]:.4f}")
                    pbar.update()

                # Evaluate after every epoch
                self.test_loss_history.append(self.predict(show_progress=False)["loss"])


    def plot_train(self, granularity=1):
        train_loss = self.train_loss_history[::granularity]
        x_iter = range(1, len(train_loss) * granularity + 1, granularity)

        n_iters = len(self.train_loss_history)
        n_epochs = len(self.test_loss_history)
        n_batches = int(n_iters / n_epochs) if n_epochs != 0 else 1
        test_loss = self.test_loss_history
        x_epoch = range(n_batches, n_iters + 1, n_batches)

        fig, axes = plt.subplots(1, 2, figsize=(13, 5))
        fig.suptitle(f"Loss and Performance Metrics During Training")

        # Left graph
        ax0 = axes[0]
        ax0.set_title("Loss")
        ax0.plot(x_iter, train_loss, label="train")
        ax0.plot(x_epoch, test_loss, label="test")
        ax0.set_xlabel("iteration")
        ax0.set_ylabel("loss")
        ax0.legend()

        # Right graph
        ax1 = axes[1]
        ax1.set_title("Training Performance")
        if len(self.train_metrics_history) > 0:
            for k in self.train_metrics_history[0].keys():
                ax1.plot(x_iter, [item[k] for item in self.train_metrics_history[::granularity]], label=k)
        ax1.set_xlabel("iteration")
        ax1.set_ylabel("score")
        ax1.set_ylim(0, 1)
        if ax1.get_legend_handles_labels()[0]:
            ax1.legend()

        return fig, axes


    def predict(self, *, show_progress=True):
        self.test_results = {"loss": 0, "probabilities": [], "predictions": []}
        all_labels = []

        self.model.eval()  # Set model to evaluation mode
        with tqdm(range(len(self.test_loader)), desc="Testing", ncols=150, disable=not show_progress) as pbar:
            with torch.no_grad():
                for batch_images, batch_labels in self.test_loader:
                    batch_images, batch_labels = batch_images.to(self.device), batch_labels.float().reshape(-1, 1).to(self.device)

                    # Predict
                    noise = v2.GaussianNoise(sigma=2)
                    batch_images = noise(batch_images)
                    batch_probas = self.model(batch_images)
                    batch_preds = batch_probas >= 0.5

                    # Record
                    all_labels.extend(batch_labels.cpu().numpy())
                    self.test_results["probabilities"].extend(batch_probas.cpu().numpy().ravel().tolist())
                    self.test_results["predictions"].extend(batch_preds.cpu().numpy().astype(np.int8).ravel().tolist())
                    self.test_results["loss"] += self.criterion(batch_probas, batch_labels).item()
                    pbar.update()

        self.test_results["metrics"] = self.__measure_metrics(all_labels, self.test_results["predictions"])
        return self.test_results
# ### Performance without training
torch.cuda.is_available()
!pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
clf = BinaryImageClassifier(lr=0.001)
clf.fit_kfold(n_epochs=10)
clf = BinaryImageClassifier(lr=0.001)
# fara noise (0.16597137786448002,
# {'accuracy': 0.98, 'precision': 0.9897959183673469, 'recall': 0.97})

out = clf.predict()
out["loss"], out["metrics"]
# ## Finetuning
# Finetune this ResNet using data in `data/train`. Use `data/test` as your testing set, and use cross entropy loss. The rest is up to you. Run finetuning that terminates within approx 10 mins. Store the following info every 10 minibatches: loss, precision, recall and accuracy on train and test datasets.
# Freeze all parameters except those in the fc layer
for param in clf.model.parameters():
    param.requires_grad = False

for param in clf.model.fc.parameters():
    param.requires_grad = True
# > The following cell takes ~3 mins to run on a T4 GPU.
clf.fit(n_epochs=10)
clf.plot_train(granularity=10);
out = clf.predict()
out["loss"], out["metrics"]
torch.save(clf.model.state_dict(), 'resnet34_finetuned.pth')
# Training set -> KFold CV -> 5 splituri pentru setul de antreanre 400/100
# finetunati de 5 ori
# augumentare -> augmix
# ## Writeup
# Summarize what you did above, as well as detail the choices you made and why. What was the outcome?
# ---
# 
# A ResNet34 classifies [1000 classes](https://deeplearning.cms.waikato.ac.nz/user-guide/class-maps/IMAGENET/), one of which is chihuahua. Likely because of this, we achieved a perfect score on the test set with minimal fine-tuning.
# 
# We loaded a pretrained ResNet34 model and replaced its final fully connected (FC) layer so that it produces a single output. This output is then passed through a sigmoid activation to clamp the value between 0 and 1.
# 
# We frozen all other layers and trained only the new FC layer. In 10 epochs, both the training and test loss decreased steadily, and finally the model reached 100% accuracy.
# 
# Before moving on, I would like to discuss some mistakes I made while making the code deterministic. I spent a lot of time on repetitive tasks as the hyperparameters obtained did not reproduce the same results, and often the performance was very different!
# 
# 1. I used to call `self.__set_seed()` only in the `__init__` method. However, during the tuning process, I found that the results were not the same each time I called `fit()`. I then realized that I should set the seed every time before training and testing, not just once during initialization.
#    
# 2. When using `DataLoader(dataset, batch_size)`, I was getting batches with all zeros followed by ones. I wanted to achieve a mix of zeros and ones in each batch (which is necessary for the next section). My initial attempt was to use `shuffle=True`, but I later realized that it shuffles the dataset at the start of each epoch, causing different runs to produce different results. Eventually, I came up with a solution by writing a custom sampler and manually shuffling the dataset.
# 
#    ```python
#     class FixedOrderSampler(Sampler):
#     ...
# 
#     class BinaryImageClassifier:
#         ...
#         def load_data(self, ...):
#             ...
#             indices = list(range(len(dataset)))
#             random.shuffle(indices)
#             indices = indices[start:stop:step]
#             
#             loader = DataLoader(dataset=dataset, batch_size=64, sampler=FixedOrderSampler(indices))
#         ...
#     ```
# 
# However, the behavior still differs between Colab and my local Jupyter Notebook for reasons that are unknown to me.
# ## Further analysis
# Pick one aspect about the work done above thus far that you find interesting, investigate it a bit further, and give a short paragraph writeup of what you investigated and how it went.
# ---
# 
# ### Defining the problem
# 
# The original problem was relatively easy to train. To make it more challenging, we impose the following restrictions:
# 
# - Limit the training set to 400 images.
# - Remove pretrained weights (i.e., start ResNet from scratch).
# 
# Our objective is to maximize classification performance (the accuracy, precision, and recall) under these constraints by tuning only three hyperparameters:
# 
# - Learning rate (`lr`)
# - Number of epochs (`n_epochs`)
# - L2 penalty (`l2_penalty`), which corresponds to `weight_decay` in `torch.optim.Adam`
# 
# We do not consider any other modifications beyond these adjustments.
# ### Setup
clfx = BinaryImageClassifier(lr=0.001, pretrained=False)
clfx.train_loader = clfx.load_data("./train", 0, 400)
clfx.test_loader = clfx.load_data("./train", 400, 500)

len(clfx.train_loader.sampler), len(clfx.test_loader.sampler)
# Here comes the significance of shuffling the images (as mentioned in the "Writeup" section). We don't want the model to train on batches where all images are labeled as zeros.
# 
# Besides, we set aside 100 images from the `./train` folder as the test set. A larger test set would provide more stable and reliable performance estimates.
# ### Tuning
# ### Learning rate
# > The following cell takes ~15 mins to run on a T4 GPU.
his = []
for lr_exp in range(-5, 0):
    clfx.l2_penalty = 0
    clfx.lr = 10 ** lr_exp
    print(f"Experimenting lr={clfx.lr}...")

    clfx.fit(20)
    his.append((clfx.lr, clfx.train_loss_history, clfx.test_loss_history))
granularity = 10
fig, axes = plt.subplots(1, 2, figsize=(13, 5))
fig.suptitle("Performance on Different Learning Rates")

# Left graph
ax0 = axes[0]
for lr, train_loss, _ in his:
    ax0.plot(range(1, len(train_loss[::granularity]) * granularity + 1, granularity), train_loss[::granularity], label=f"lr={lr}")
ax0.set_xlabel("iteration")
ax0.set_ylabel("loss")
ax0.set_title("Train Loss")
ax0.set_ylim(0, 1)
ax0.legend()

# Right graph
ax1 = axes[1]
for lr, _, test_loss in his:
    ax1.plot(range(1, len(test_loss) + 1), test_loss, label=f"lr={lr}")
ax1.set_xlabel("epoch")
ax1.set_ylabel("loss")
ax1.set_title("Test Loss")
ax1.set_ylim(0, 5)
ax1.legend();
# ### Number of epochs
# Using the best `lr` found from the graph above, we now train for 50 epochs to take a closer look on how the test loss evolves over time.
# > The following cell takes ~7 mins to run on a T4 GPU.
clfx.lr = 0.0001
clfx.l2_penalty = 0
clfx.fit(50)
fig, axes = clfx.plot_train(10)

min_idx = np.argmin(clfx.test_loss_history)
min_x = (min_idx + 1) * len(clfx.train_loader)
min_y = clfx.test_loss_history[min_idx]

axes[0].scatter(min_x, min_y, color='red')
axes[0].annotate(f'min loss: {min_y:.4f}\n(n_epochs={min_idx + 1})', (min_x, min_y), textcoords="offset points", xytext=(5,7));
# ### L2 penalty
# After some empirical testing, it appears that adding an L2 penalty did not improve the scores. Therefore, we set it to zero. Our final result is:
clfx.lr = 0.0001
clfx.l2_penalty = 0
clfx.fit(16)

out = clfx.predict()
out["loss"], out["metrics"]
torch.save(clfx.model.state_dict(), 'resnet34_manual400.pth')
# ## Auto tuning
# 
# In this section, we utilize Optuna to automatically tune our hyperparameters. Optuna automates hyperparameter tuning using Bayesian optimization. It smartly selects hyperparameter configurations based on past trial outcomes, focusing on promising regions of the search space.
!pip install optuna
import optuna
import optuna.visualization as vis
def objective(trial):
    lr = trial.suggest_float('lr', 1e-5, 1e-1, log=True)
    n_epochs = trial.suggest_int('n_epochs', 1, 20)
    l2_penalty = trial.suggest_float('l2_penalty', 1e-6, 1e-1, log=True)

    clfy = BinaryImageClassifier(lr=lr, l2_penalty=l2_penalty, pretrained=False)

    clfy.train_loader = clfy.load_data("./train", 0, 400)
    clfy.test_loader = clfy.load_data("./train", 400, 500)

    clfy.fit(n_epochs)
    out = clfy.predict()
    return *out["metrics"].values(), out["loss"]
# > The following cell takes ~1 hour to run on a T4 GPU.
study = optuna.create_study(directions=["maximize", "maximize", "maximize", "minimize"],
                            sampler=optuna.samplers.TPESampler(seed=42))
study.optimize(objective, n_trials=30)
vis.plot_optimization_history(
    study,
    target=lambda t: t.values[0],
    target_name="Accuracy"
)
vis.plot_slice(
    study,
    target=lambda t: t.values[0],
    target_name="Accuracy"
)
vis.plot_param_importances(
    study,
    target=lambda t: t.values[0],
    target_name="Accuracy"
)
study.best_trials
best_params = study.trials[19].params
best_params
study.trials[19].values
clfy = BinaryImageClassifier(lr=best_params["lr"], l2_penalty=best_params["l2_penalty"], pretrained=False)

clfy.train_loader = clfy.load_data("./train", 0, 400)
clfy.test_loader = clfy.load_data("./train", 400, 500)

clfy.fit(best_params["n_epochs"])
clfy.plot_train(10);
out = clfy.predict()
out["loss"], out["metrics"]
torch.save(clfy.model.state_dict(), 'resnet34_auto400.pth')
# ## Conclusion
# 
# | | Parameters | Accuracy | Precision | Recall | Loss |
# | --- | --- | --- | --- | --- | --- |
# | Manual | `lr=1.000e-4, n_epochs=16, l2_penalty=0.000` | 0.9300 | 0.9375 | 0.9184 | 0.4191 |
# | Optuna | `lr=2.351e-4, n_epochs=20, l2_penalty=1.917e-4` | 0.9700 | 0.9792 | 0.9592 | 0.3445 |
# 
# (Correct to 4 sig. fig.)
# 
# The results demonstrate that the Optuna-optimized model outperformed the manually tuned counterpart across all evaluation metrics, achieving higher accuracy, precision, and recall, while also reducing the loss.
# 
# This highlights the effectiveness of automated hyperparameter optimization in improving model performance and reducing the need for extensive manual experimentation. Also, the manual approach of "tuning one parameter at a time" may be suboptimal, as hyperparameters often interact in non-orthogonal ways: changes in one can influence the optimal values of others.
