하이퍼파라미터

import torch
import torch.nn as nn
import torch.optim as optim
from dataclasses import dataclass
from typing import Optional

@dataclass
class HyperParameters:
    """모델 학습을 위한 하이퍼파라미터 설정"""
    # 모델 구조 하이퍼파라미터
    hidden_dim: int = 256
    num_layers: int = 3
    num_heads: int = 8  # Transformer용
    dropout: float = 0.1

    # 학습 알고리즘 하이퍼파라미터
    learning_rate: float = 1e-4
    weight_decay: float = 0.01
    batch_size: int = 32
    num_epochs: int = 100

    # 옵티마이저 하이퍼파라미터
    optimizer: str = "adamw"  # adam, sgd, adamw
    beta1: float = 0.9
    beta2: float = 0.999
    momentum: float = 0.9  # SGD용

    # 학습률 스케줄러 하이퍼파라미터
    scheduler: str = "cosine"  # cosine, step, linear
    warmup_steps: int = 1000
    min_lr: float = 1e-6

    # 정규화 하이퍼파라미터
    gradient_clip: float = 1.0
    early_stopping_patience: int = 10
    label_smoothing: float = 0.1

    # LoRA 파인튜닝 하이퍼파라미터 (LLM용)
    lora_rank: int = 8
    lora_alpha: int = 16
    lora_dropout: float = 0.05

class SimpleMLP(nn.Module):
    """하이퍼파라미터 기반 MLP 모델"""
    def __init__(self, input_dim: int, output_dim: int, hp: HyperParameters):
        super().__init__()
        layers = []
        dims = [input_dim] + [hp.hidden_dim] * hp.num_layers + [output_dim]

        for i in range(len(dims) - 1):
            layers.append(nn.Linear(dims[i], dims[i+1]))
            if i < len(dims) - 2:  # 마지막 층 제외
                layers.append(nn.ReLU())
                layers.append(nn.Dropout(hp.dropout))

        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x)

def create_optimizer(model: nn.Module, hp: HyperParameters) -> optim.Optimizer:
    """하이퍼파라미터에 따른 옵티마이저 생성"""
    if hp.optimizer == "adam":
        return optim.Adam(
            model.parameters(),
            lr=hp.learning_rate,
            betas=(hp.beta1, hp.beta2),
            weight_decay=hp.weight_decay
        )
    elif hp.optimizer == "adamw":
        return optim.AdamW(
            model.parameters(),
            lr=hp.learning_rate,
            betas=(hp.beta1, hp.beta2),
            weight_decay=hp.weight_decay
        )
    elif hp.optimizer == "sgd":
        return optim.SGD(
            model.parameters(),
            lr=hp.learning_rate,
            momentum=hp.momentum,
            weight_decay=hp.weight_decay
        )
    else:
        raise ValueError(f"Unknown optimizer: {hp.optimizer}")

def create_scheduler(optimizer: optim.Optimizer, hp: HyperParameters, total_steps: int):
    """학습률 스케줄러 생성"""
    if hp.scheduler == "cosine":
        return optim.lr_scheduler.CosineAnnealingLR(
            optimizer,
            T_max=total_steps - hp.warmup_steps,
            eta_min=hp.min_lr
        )
    elif hp.scheduler == "step":
        return optim.lr_scheduler.StepLR(
            optimizer,
            step_size=total_steps // 3,
            gamma=0.1
        )
    elif hp.scheduler == "linear":
        return optim.lr_scheduler.LinearLR(
            optimizer,
            start_factor=1.0,
            end_factor=hp.min_lr / hp.learning_rate,
            total_iters=total_steps
        )

def train_with_hyperparameters(
    model: nn.Module,
    train_loader,
    hp: HyperParameters
):
    """하이퍼파라미터를 사용한 학습 루프"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    criterion = nn.CrossEntropyLoss(label_smoothing=hp.label_smoothing)
    optimizer = create_optimizer(model, hp)

    total_steps = len(train_loader) * hp.num_epochs
    scheduler = create_scheduler(optimizer, hp, total_steps)

    best_loss = float('inf')
    patience_counter = 0

    for epoch in range(hp.num_epochs):
        model.train()
        epoch_loss = 0.0

        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)

            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()

            # 그래디언트 클리핑 적용
            torch.nn.utils.clip_grad_norm_(model.parameters(), hp.gradient_clip)

            optimizer.step()
            scheduler.step()

            epoch_loss += loss.item()

        avg_loss = epoch_loss / len(train_loader)
        print(f"Epoch {epoch+1}/{hp.num_epochs}, Loss: {avg_loss:.4f}, LR: {scheduler.get_last_lr()[0]:.2e}")

        # 조기 종료 체크
        if avg_loss < best_loss:
            best_loss = avg_loss
            patience_counter = 0
        else:
            patience_counter += 1
            if patience_counter >= hp.early_stopping_patience:
                print(f"Early stopping at epoch {epoch+1}")
                break

# 사용 예시
if __name__ == "__main__":
    # 하이퍼파라미터 설정
    hp = HyperParameters(
        hidden_dim=512,
        num_layers=4,
        learning_rate=3e-4,
        batch_size=64,
        dropout=0.2
    )
    print(f"하이퍼파라미터: {hp}")

    # 모델 생성
    model = SimpleMLP(input_dim=784, output_dim=10, hp=hp)
    print(f"모델 파라미터 수: {sum(p.numel() for p in model.parameters()):,}")