model_3.py

# -*- coding: utf-8 -*-
"""Assignment#3.ipynb의 사본

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/13tX_bkgPrVxOx7MhBgBltv7zvXwZ-W6i

# **Assignment#3**

Conducting a simple experiment for multi-classification using CIFAR-10 with MLP

*   Hyperparameter tuning
    - Find the combination of hyperparameters based on your own idea

1.   No. of hidden layers, No. of hidden units
2.   Learning rate
3.   Weight initialization depending on activation function
     - (Glorot) Xavier, (Kaiming) He
4.   Regularization
     - L2
     - Dropout

Note that you can report your final results to use other frameworks, such as Tensorflow, Theano, Julia, and etc.

# **Submitting your final results about "Accuracy" in terms of training, validation and Testing based on the optimal combination of hyperparameters you set up (by 11/04/25 23:59)**

 - The report should include codes and final results (figures)
----------------------------------------------------------------------------------

## Creating a directory "results"
"""

# !mkdir results


import os, time, json, argparse, hashlib
import numpy as np
import pandas as pd
from copy import deepcopy

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Subset, Dataset, DataLoader

import torchvision
import torchvision.transforms as transforms

import matplotlib.pyplot as plt
import seaborn as sns
import optuna

"""# **Data Preparation**
- Use CIFAR10
"""

os.environ["CUDA_VISIBLE_DEVICES"] = "1"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
os.makedirs("results_optuna", exist_ok=True)

# (A) 변환 정의
train_transform = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ColorJitter(0.2, 0.2, 0.2, 0.1),
    transforms.ToTensor(),
    transforms.RandomErasing(p=0.25, scale=(0.02, 0.15), value=0.0),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])

eval_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])

# (B) Subset 에 서로 다른 transform 적용하기 위한 래퍼
class TransformedSubset(Dataset):
    def __init__(self, subset, transform):
        self.subset = subset          # torch.utils.data.Subset
        self.transform = transform
    def __len__(self):
        return len(self.subset)
    def __getitem__(self, i):
        x, y = self.subset.dataset[self.subset.indices[i]]
        x = self.transform(x)
        return x, y

# (C) 데이터셋/분할
base_train = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=None)
indices = np.arange(len(base_train))
train_idx, val_idx = indices[:40000], indices[40000:]

train_subset = Subset(base_train, train_idx)
val_subset   = Subset(base_train, val_idx)

trainset = TransformedSubset(train_subset, train_transform)
valset   = TransformedSubset(val_subset,   eval_transform)
testset  = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=eval_transform)

partition = {'train': trainset, 'val': valset, 'test': testset}

"""# **Model Architecture**"""

class MLP(nn.Module):
    def __init__(self, in_dim, out_dim, hid_dim, n_layer, act, dropout, use_bn, w_initial):
        super(MLP, self).__init__()
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.hid_dim = hid_dim
        self.n_layer = n_layer
        self.act = act
        self.dropout = dropout
        self.use_bn = use_bn
        self.w_initial = w_initial

        # ====== Create Linear Layers ====== #
        self.fc1 = nn.Linear(self.in_dim, self.hid_dim)

        self.linears = nn.ModuleList()
        self.bns = nn.ModuleList()
        for i in range(self.n_layer-1):
            self.linears.append(nn.Linear(self.hid_dim, self.hid_dim))
            if self.use_bn:
                self.bns.append(nn.BatchNorm1d(self.hid_dim))

        self.fc2 = nn.Linear(self.hid_dim, self.out_dim)

        # ====== Create Activation Function ====== #
        if self.act == 'relu':
            self.act = nn.ReLU()
            self._kaiming_nl = 'relu'
        elif self.act == 'leakyrelu':
            self.act = nn.LeakyReLU()
            self._kaiming_nl = 'leaky_relu'
        elif self.act == 'tanh':
            self.act = nn.Tanh()
            self._kaiming_nl = 'tanh'
        elif self.act == 'sigmoid':
            self.act = nn.Sigmoid()
            self._kaiming_nl = 'sigmoid'
        else:
            raise ValueError('no valid activation function selected!')

        # ====== Create Regularization Layer ======= #
        self.dropout = nn.Dropout(self.dropout)

        if self.w_initial == 'xavier':
            self.xavier_init()
        elif self.w_initial == 'he':
            self.he_init()
        else:
            raise ValueError("no valid initialization method selected!")

    def forward(self, x):
        x = self.act(self.fc1(x))
        for i in range(len(self.linears)):
            x = self.linears[i](x)
            if self.use_bn:
                x = self.bns[i](x)            # Linear -> BN -> Act -> Dropout
            x = self.act(x)
            x = self.dropout(x)
        x = self.fc2(x)
        return x

    # ---- 초기화에 fc1/fc2도 포함하도록 수정 ----
    def xavier_init(self):
        nn.init.xavier_normal_(self.fc1.weight)
        self.fc1.bias.data.fill_(0.01)
        for linear in self.linears:
            nn.init.xavier_normal_(linear.weight)
            linear.bias.data.fill_(0.01)
        nn.init.xavier_normal_(self.fc2.weight)
        self.fc2.bias.data.zero_()

    def he_init(self):
        nl = getattr(self, "_kaiming_nl", "leaky_relu")
        nn.init.kaiming_normal_(self.fc1.weight, nonlinearity=nl)
        self.fc1.bias.data.fill_(0.01)
        for linear in self.linears:
            nn.init.kaiming_normal_(linear.weight, nonlinearity=nl)
            linear.bias.data.fill_(0.01)
        nn.init.kaiming_normal_(self.fc2.weight, nonlinearity=nl)
        self.fc2.bias.data.zero_()


net = MLP(3072, 10, 100, 4, 'tanh', 0.1, True, 'he') # in_dim, out_dim, hid_dim, n_layer, act, dropout, use_bn, w_initial
net

"""# **train, validation, test and experiment**"""

def train(net, partition, optimizer, criterion, args):
    trainloader = torch.utils.data.DataLoader(partition['train'],
                                              batch_size=args.train_batch_size,
                                              shuffle=True, num_workers=2)
    net.train()

    correct = 0
    total = 0
    train_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        optimizer.zero_grad()
        # get the inputs
        inputs, labels = data
        inputs = inputs.view(-1, 3072)
        inputs = inputs.cuda()
        labels = labels.cuda()
        outputs = net(inputs)

        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    train_loss = train_loss / len(trainloader)
    train_acc = 100 * correct / total
    return net, train_loss, train_acc

def validate(net, partition, criterion, args):
    valloader = torch.utils.data.DataLoader(partition['val'],
                                            batch_size=args.test_batch_size,
                                            shuffle=False, num_workers=2)
    net.eval()

    correct = 0
    total = 0
    val_loss = 0
    with torch.no_grad():
        for data in valloader:
            images, labels = data
            images = images.view(-1, 3072)
            images = images.cuda()
            labels = labels.cuda()
            outputs = net(images)

            loss = criterion(outputs, labels)

            val_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        val_loss = val_loss / len(valloader)
        val_acc = 100 * correct / total
    return val_loss, val_acc

def test(net, partition, args):
    testloader = torch.utils.data.DataLoader(partition['test'],
                                             batch_size=args.test_batch_size,
                                             shuffle=False, num_workers=2)
    net.eval()

    correct = 0
    total = 0
    with torch.no_grad():
        for data in testloader:
            images, labels = data
            images = images.view(-1, 3072)
            images = images.cuda()
            labels = labels.cuda()

            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        test_acc = 100 * correct / total
    return test_acc

def experiment(partition, args, trial=None):
    # net = MLP(3072, 10, 100, 4, 'relu', 0.1, True, 'he') # in_dim, out_dim, hid_dim, n_layer, act, dropout, use_bn, w_initial
    net = MLP(args.in_dim, args.out_dim, args.hid_dim, args.n_layer, args.act, args.dropout, args.use_bn, args.w_initial)
    net.cuda()

    criterion = nn.CrossEntropyLoss(label_smoothing=0.05)
    if args.optim == 'SGD':
        optimizer = optim.SGD(net.parameters(), lr=args.lr, weight_decay=args.l2)
    elif args.optim == 'RMSprop':
        optimizer = optim.RMSprop(net.parameters(), lr=args.lr, weight_decay=args.l2)
    elif args.optim == 'ADAM':
        optimizer = optim.Adam(net.parameters(), lr=args.lr, weight_decay=args.l2)
    elif args.optim == 'AdamW':
        optimizer = optim.AdamW(net.parameters(), lr=args.lr, weight_decay=args.l2)  # <-- AdamW로 실제 적용
    else:
        raise ValueError('In-valid optimizer choice')

    train_losses = []
    val_losses = []
    train_accs = []
    val_accs = []

    best_val_loss = float('inf')
    best_val_acc = 0.0        # best validation accuracy
    best_state = None         # best paramter
    patience = 5              # patience
    wait = 0                  # 개선이 없었던 epoch 수
    min_delta = 1e-4

    for epoch in range(args.epoch):
        ts = time.time()
        net, train_loss, train_acc = train(net, partition, optimizer, criterion, args)
        val_loss, val_acc = validate(net, partition, criterion, args)     
        te = time.time()

        train_losses.append(train_loss) 
        val_losses.append(val_loss)
        train_accs.append(train_acc)
        val_accs.append(val_acc)

        print(f'Epoch {epoch}, Acc(train/val): {train_acc:.2f}/{val_acc:.2f}, '
              f'Loss(train/val): {train_loss:.2f}/{val_loss:.2f}, Took {te-ts:.2f}s')

        # ---- (NEW) Optuna pruning hook ----
        if trial is not None:
            trial.report(val_acc, step=epoch)
            if trial.should_prune():
                raise optuna.TrialPruned()

        # --- 조기종료 체크 (Validation Loss 단일 기준으로 안정화) ---
        if val_loss < best_val_loss - min_delta:
            best_val_loss = val_loss
            best_val_acc = val_acc        # 같은 시점의 val_acc도 저장
            best_state = deepcopy(net.state_dict())
            wait = 0
        else:
            wait += 1
            if wait >= patience:
                print(f"[Early Stop] No improvement in loss for {patience} epochs. Stop.")
                break

    if best_state is not None:
        net.load_state_dict(best_state)
        print(f'[Final Result] Best Val Acc: {best_val_acc:.2f}, Best Val Loss: {best_val_loss:.2f}')

    test_acc = test(net, partition, args)

    result = {}
    result['train_losses'] = train_losses
    result['val_losses'] = val_losses
    result['train_accs'] = train_accs
    result['val_accs'] = val_accs
    result['train_acc'] = max(train_accs) if train_accs else 0.0  # 보고용: 최고 train acc
    result['val_acc'] = best_val_acc                              # best loss 시점의 val acc
    result['test_acc'] = test_acc
    return vars(args), result

"""# **Managing experimental results**
"""


import hashlib
import json
from os import listdir
from os.path import isfile, join
import pandas as pd

def save_exp_result(setting, result):
    exp_name = setting['exp_name']
    del setting['epoch']
    del setting['test_batch_size']

    hash_key = hashlib.sha1(str(setting).encode()).hexdigest()[:6]
    filename = './results_optuna/{}-{}.json'.format(exp_name, hash_key)
    result.update(setting)
    with open(filename, 'w') as f:
        json.dump(result, f)


def load_exp_result(exp_name):
    dir_path = './results_optuna'
    filenames = [f for f in listdir(dir_path) if isfile(join(dir_path, f)) if '.json' in f]
    list_result = []
    for filename in filenames:
        if exp_name in filename:
            with open(join(dir_path, filename), 'r') as infile:
                results = json.load(infile)
                list_result.append(results)
    df = pd.DataFrame(list_result) # .drop(columns=[])
    return df

"""# **Visualizing the result figures**"""

def plot_acc(var1, var2, df):

    fig, ax = plt.subplots(1, 3)
    fig.set_size_inches(15, 6)
    sns.set_style("darkgrid", {"axes.facecolor": ".9"})

    sns.barplot(x=var1, y='train_acc', hue=var2, data=df, ax=ax[0])
    sns.barplot(x=var1, y='val_acc', hue=var2, data=df, ax=ax[1])
    sns.barplot(x=var1, y='test_acc', hue=var2, data=df, ax=ax[2])

    ax[0].set_title('Train Accuracy')
    ax[1].set_title('Validation Accuracy')
    ax[2].set_title('Test Accuracy')


def plot_loss_variation(var1, var2, df, **kwargs):

    list_v1 = df[var1].unique()
    list_v2 = df[var2].unique()
    list_data = []

    for value1 in list_v1:
        for value2 in list_v2:
            row = df.loc[df[var1]==value1]
            row = row.loc[row[var2]==value2] 

            train_losses = list(row.train_losses)[0]
            val_losses = list(row.val_losses)[0]

            for epoch, train_loss in enumerate(train_losses):
                list_data.append({'type':'train', 'loss':train_loss, 'epoch':epoch, var1:value1, var2:value2})
            for epoch, val_loss in enumerate(val_losses):
                list_data.append({'type':'val', 'loss':val_loss, 'epoch':epoch, var1:value1, var2:value2})

    df = pd.DataFrame(list_data)
    g = sns.FacetGrid(df, row=var2, col=var1, hue='type', height=4, aspect=1.2, **kwargs)
    g = g.map(plt.plot, 'epoch', 'loss', marker='.')
    g.add_legend()
    g.figure.suptitle('Train loss vs Val loss', fontsize=16)
    plt.tight_layout(rect=[0, 0.03, 1, 0.95])


def plot_acc_variation(var1, var2, df, **kwargs):
    list_v1 = df[var1].unique()
    list_v2 = df[var2].unique()
    list_data = []

    for value1 in list_v1:
        for value2 in list_v2:
            row = df.loc[df[var1]==value1]
            row = row.loc[row[var2]==value2]   # <-- 버그 수정

            train_accs = list(row.train_accs)[0]
            val_accs = list(row.val_accs)[0]
            test_acc = list(row.test_acc)[0]

            for epoch, train_acc in enumerate(train_accs):
                list_data.append({'type':'train', 'Acc':train_acc, 'test_acc':test_acc, 'epoch':epoch, var1:value1, var2:value2})
            for epoch, val_acc in enumerate(val_accs):
                list_data.append({'type':'val', 'Acc':val_acc, 'test_acc':test_acc, 'epoch':epoch, var1:value1, var2:value2})

    df = pd.DataFrame(list_data)
    g = sns.FacetGrid(df, row=var2, col=var1, hue='type', height=4, aspect=1.2, **kwargs)
    g = g.map(plt.plot, 'epoch', 'Acc', marker='.')

    def show_acc(x, y, metric, **kwargs):
        plt.scatter(x, y, alpha=0.3, s=1)
        metric = "Test Acc: {:1.3f}".format(list(metric.values)[0])
        plt.text(0.05, 0.95, metric,  horizontalalignment='left', verticalalignment='center', transform=plt.gca().transAxes, bbox=dict(facecolor='yellow', alpha=0.5, boxstyle="round,pad=0.1"))
    g = g.map(show_acc, 'epoch', 'Acc', 'test_acc')

    g.add_legend()
    g.figure.suptitle('Train Accuracy vs Val Accuracy', fontsize=16) 
    plt.tight_layout(rect=[0, 0.03, 1, 0.95]) 

"""# **Optuna**
"""

def run_optuna(partition, base_args, n_trials=60, search_epochs=30, final_epochs=50, seed=42):
    # Sampler / Pruner
    sampler = optuna.samplers.TPESampler(multivariate=True, group=True, n_startup_trials=12, seed=seed)
    pruner  = optuna.pruners.SuccessiveHalvingPruner(min_resource=5, reduction_factor=3, min_early_stopping_rate=0)

    def objective(trial: optuna.Trial):
        args = deepcopy(base_args)

        # ---- 기존 하이퍼파라미터만 탐색 ----
        args.hid_dim = trial.suggest_categorical("hid_dim", [128, 192, 256, 320, 384, 512])
        args.n_layer = trial.suggest_int("n_layer", 3, 6)
        args.dropout = trial.suggest_float("dropout", 0.0, 0.6)
        args.use_bn  = trial.suggest_categorical("use_bn", [True, False])
        args.lr      = trial.suggest_float("lr", 5e-4, 1e-2, log=True)
        args.l2      = trial.suggest_float("l2", 1e-6, 3e-3, log=True)
        args.act     = trial.suggest_categorical("act", ["relu", "leakyrelu", "tanh"])
        # init은 act에 조건부로 매핑
        args.w_initial = "he" if args.act in ["relu", "leakyrelu"] else "xavier"
        args.train_batch_size = trial.suggest_categorical("train_batch_size", [256, 512, 768, 1024])

        # 고정
        args.optim = 'AdamW'
        args.epoch = search_epochs  # 탐색은 짧게
        # 실행
        _, result = experiment(partition, args, trial=trial)
        return result['val_acc']

    study = optuna.create_study(direction="maximize", sampler=sampler, pruner=pruner)
    study.optimize(objective, n_trials=n_trials, show_progress_bar=True)

    print("\n==== Optuna Best ====")
    print("Best value (val_acc):", study.best_value)
    print("Best params:", study.best_params)

    # ----- 베스트 설정으로 재학습 (곡선/요약 = '베스트만' 저장) -----
    best = deepcopy(base_args)
    best.hid_dim = study.best_params['hid_dim']
    best.n_layer = study.best_params['n_layer']
    best.dropout = study.best_params['dropout']
    best.use_bn  = study.best_params['use_bn']
    best.lr      = study.best_params['lr']
    best.l2      = study.best_params['l2']
    best.act     = study.best_params['act']
    best.w_initial = "he" if best.act in ["relu", "leakyrelu"] else "xavier"
    best.train_batch_size = study.best_params['train_batch_size']

    best.optim = 'AdamW'
    best.epoch = final_epochs

    # 재학습 (pruning 없음)
    setting, result = experiment(partition, best, trial=None)

    # ---- 결과 저장: results_optuna/ 베스트만 ----
    # 곡선
    plt.figure(figsize=(7,5))
    plt.plot(result['train_losses'], label="Train")
    plt.plot(result['val_losses'], label="Val")
    plt.xlabel("Epoch"); plt.ylabel("Loss"); plt.title("Best (Optuna) - Loss")
    plt.legend(); plt.tight_layout()
    plt.savefig("results_optuna/best_loss_curve.png", dpi=150)

    plt.figure(figsize=(7,5))
    plt.plot(result['train_accs'], label="Train")
    plt.plot(result['val_accs'], label="Val")
    plt.xlabel("Epoch"); plt.ylabel("Accuracy (%)"); plt.title("Best (Optuna) - Accuracy")
    plt.legend(); plt.tight_layout()
    plt.savefig("results_optuna/best_acc_curve.png", dpi=150)

    # 요약/설정 저장
    with open("results_optuna/best_summary.txt", "w") as f:
        json.dump({
            "best_val_acc": float(result['val_acc']),
            "test_acc": float(result['test_acc']),
            "train_best_acc": float(result['train_acc']),
            "epochs_run": int(len(result['train_losses']))
        }, f, indent=2)

    # base_args/최종 best 설정 저장
    best_config = deepcopy(vars(best))
    with open("results_optuna/best_config.json", "w") as f:
        json.dump(best_config, f, indent=2)

    print("\n[Optuna Saved] results_optuna/: best_loss_curve.png, best_acc_curve.png, best_summary.txt, best_config.json")


"""# **Experiment No.1 : No. of hidden layers and hidden units**"""

# ====== Random Seed Initialization ====== #
seed = 42
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

# ---- (UPDATED) CLI 옵션: 기본 실행이 optuna + 지정 파라미터 ----
parser = argparse.ArgumentParser()
parser.add_argument("--mode", type=str, default="optuna", choices=["grid", "optuna"])
parser.add_argument("--n_trials", type=int, default=80)
parser.add_argument("--search_epochs", type=int, default=30)
parser.add_argument("--final_epochs", type=int, default=50)
# Colab/Notebook 호환: 노트북이면 args=[], 스크립트면 실제 argv 사용
args_cli = parser.parse_args(args=[] if '__file__' not in globals() else None)

# (원본 args 유지)
args = argparse.Namespace()
args.exp_name = "exp2_lr"

# 층 수
args.n_layer = 4 
# 은닉 유닛 수
args.hid_dim = 256 

# ====== Model Capacity ====== #
args.in_dim = 3072
args.out_dim = 10
args.act = 'leakyrelu' # leakyrelu

# ====== Regularization ======= #
args.dropout = 0.4
args.use_bn = True
args.l2 = 0.00001
args.w_initial = 'he' # he

# ====== Optimizer & Training ====== #
args.optim = 'AdamW'  # 'RMSprop' #SGD, RMSprop, ADAM...
args.lr = 0.001
args.epoch = 50

args.train_batch_size = 1024
args.test_batch_size = 1024

# ====== Experiment Variable (Grid) ====== #
name_var1 = 'l2'
name_var2 = 'hid_dim'
list_var1 = [1e-5, 3e-5, 1e-4, 3e-4, 1e-3]
list_var2 = [args.hid_dim]

if args_cli.mode == "grid":
    for var1 in list_var1:
        for var2 in list_var2:
            setattr(args, name_var1, var1)
            setattr(args, name_var2, var2)
            print(args)

            setting, result = experiment(partition, deepcopy(args))
            save_exp_result(setting, result)

    # ---- 시각화/출력 ----
    var1 = 'l2'
    var2 = 'hid_dim'
    df_lr = load_exp_result('exp2_lr')

    plot_loss_variation(var1, var2, df_lr, sharey=False) 
    plt.savefig("results_optuna/loss_variation_l2_tuning.png")
    plt.show()

    plot_acc_variation(var1, var2, df_lr, margin_titles=True, sharey=True) 
    plt.savefig("results_optuna/accuracy_variation_l2_tuning.png")
    plt.show()

    print("\n=== Exp2 (L2) 최종 결과 테이블 ===")
    print(df_lr[['l2', 'train_acc', 'val_acc', 'test_acc']].sort_values(by='val_acc', ascending=False))
    print("==================================")

else:
    # ---- Optuna 실행: 기본 실행이 여기로 진입 ----
    base_args = deepcopy(args)
    run_optuna(
        partition=partition,
        base_args=base_args,
        n_trials=args_cli.n_trials,        
        search_epochs=args_cli.search_epochs,  
        final_epochs=args_cli.final_epochs,   
        seed=seed
    )