PSO/example/pso2mnist.ipynb

pso2keras PyTorch 3.2.0 MNIST PSO Optimization

이 노트북은 PyTorch nn.Module과 pso2keras (v3.2.0) 라이브러리를 사용하여 MNIST 이미지 분류 모델을 Particle Swarm Optimization (PSO) 알고리즘으로 최적화하는 예제입니다.

환경 요구사항:

• Python 3.11 및 PyTorch >= 2.13.0
• uv 패키지 관리자를 사용한 pso2keras[examples] 설치
• MNIST 데이터셋 최초 다운로드를 위한 인터넷 연결 필요 (torchvision.datasets.MNIST)

0. 패키지 설치 (Google Colab 및 시스템 환경)

uv 패키지 관리자를 설치하고 pso2keras[examples] 패키지를 시스템 파이썬 환경에 설치합니다.

In [ ]:

# Google Colab 및 시스템 환경 uv 패키지 설치
!pip install -q uv
!uv pip install 'pso2keras[examples]' --system

In [ ]:

import sys
import torch

print("Python version:", sys.version)
print("PyTorch version:", torch.__version__)

# Metal MPS (Apple Silicon GPU) 가속 백엔드 진단
built = hasattr(torch.backends, "mps") and torch.backends.mps.is_built()
avail = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
print(f"MPS Backend - Built: {built}, Available: {avail}")

1. 데이터셋 다운로드 및 PCA 전처리

torchvision.datasets.MNIST를 이용해 데이터셋을 다운로드하고, 결정론적(Deterministic) 서브셋을 추출합니다. PSO 알고리즘의 파티클 탐색 효율을 높이기 위해 scikit-learn PCA를 사용하여 28x28 (784차원) 이미지를 50차원 피처 표현으로 압축합니다. 다중 클래스 분류(task="multiclass")를 위해 타겟 레이블은 1D 정수형(torch.long / int64) 클래스 인덱스로 구성합니다.

In [ ]:

import torch
import torch.nn as nn
from torchvision import datasets, transforms
from sklearn.decomposition import PCA
from pso import Optimizer

# 재현 가능한 시드 설정
torch.manual_seed(42)

# 1. MNIST 데이터셋 로드 (인터넷 다운로드 필요)
transform = transforms.Compose([transforms.ToTensor()])
mnist_train = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
mnist_test = datasets.MNIST(root="./data", train=False, download=True, transform=transform)

# 2. 결정론적 샘플 서브셋 추출 (학습 5,000개, 검증 1,000개)
g = torch.Generator().manual_seed(42)
train_indices = torch.randperm(len(mnist_train), generator=g)[:5000]
test_indices = torch.randperm(len(mnist_test), generator=g)[:1000]

x_train_raw = mnist_train.data[train_indices].float() / 255.0
y_train = mnist_train.targets[train_indices].long()

x_test_raw = mnist_test.data[test_indices].float() / 255.0
y_test = mnist_test.targets[test_indices].long()

# 3. PCA 50차원 주성분 분석 전처리
x_train_flat = x_train_raw.view(x_train_raw.size(0), -1).numpy()
x_test_flat = x_test_raw.view(x_test_raw.size(0), -1).numpy()

pca = PCA(n_components=50, random_state=42)
x_train_pca = pca.fit_transform(x_train_flat)
x_test_pca = pca.transform(x_test_flat)

x_train = torch.tensor(x_train_pca, dtype=torch.float32)
x_test = torch.tensor(x_test_pca, dtype=torch.float32)

print(f"x_train : {x_train.shape} | y_train : {y_train.shape}")
print(f"x_test  : {x_test.shape}  | y_test  : {y_test.shape}")

2. PyTorch 신경망 모델 및 손실 함수 정의

50개의 PCA 입력 피처를 받아 10개 클래스의 Raw Logits를 출력하는 소형 신경망 MNISTLogitNet을 정의합니다. 다중 클래스 분류를 위해 nn.CrossEntropyLoss()를 사용합니다.

In [ ]:

class MNISTLogitNet(nn.Module):
    def __init__(self, in_features=50, hidden_dim=32, num_classes=10):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_features, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, num_classes)
        )

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

model = MNISTLogitNet()
loss_fn = nn.CrossEntropyLoss()

3. PSO Optimizer 생성 및 최적화 실행

Optimizer 생성자에 PyTorch 모델, 손실 함수, 작업 유형(task="multiclass"), 이동 방식(method="inertia"), 초기화 방식(initialization="model_noise"), 적합도 평가 방식(evaluation="fixed_subset"), 수렴 방식(convergence="none"), 후속 정제 방식(refinement="adam"), 파티클 개수(n_particles), 이동 계수(c0, c1, w_min, w_max), 가중치 경계(particle_min, particle_max), 속도 제한 비율(velocity_limit_ratio), 경계 전략(boundary_strategy), 초기 노이즈 스케일(initial_position_noise), 시드(seed)를 설정합니다.

device=None을 지정하면 MPS (Apple Silicon GPU) -> CUDA -> CPU 순서로 실행 디바이스를 자동 선택합니다.

fit 메서드에 고정 적합도 서브셋 크기(fitness_size=2000), 배치 크기(batch_size=500), 검증 데이터셋(validation_data=(x_test, y_test)), 하이브리드 Adam 국소 정제(refinement_epochs=100, refinement_lr=0.01), 결과 디렉토리(output_dir="./result/mnist")를 지정하여 최적화를 수행합니다.

In [ ]:

pso_mnist = Optimizer(
    model,
    loss=loss_fn,
    task="multiclass",
    method="inertia",
    initialization="model_noise",
    evaluation="fixed_subset",
    convergence="none",
    refinement="adam",
    fitness_size=2000,
    n_particles=20,
    c0=0.35,
    c1=0.8,
    w_min=0.6,
    w_max=1.2,
    particle_min=-3.0,
    particle_max=3.0,
    velocity_limit_ratio=0.1,
    boundary_strategy="reflect",
    initial_position_noise=0.05,
    seed=42,
    device=None,
    refinement_epochs=100,
    refinement_lr=0.01,
)

best_score = pso_mnist.fit(
    x_train,
    y_train,
    epochs=15,
    batch_size=500,
    fitness_size=2000,
    renewal="acc",
    validation_data=(x_test, y_test),
    output_dir="./result/mnist",
    log_format="csv",
    checkpoint_interval=5,
    save_info=True,
    refinement_epochs=100,
    refinement_lr=0.01,
)

print(f"Optimization Completed! Best Training Score (loss, accuracy, mse): {best_score}")
val_score = pso_mnist.evaluate(x_test, y_test)
print(f"Validation Score (loss, accuracy, mse): {val_score}")

4. 최적 모델 체크포인트 로드 및 이식 가능한 모델 검증

학습이 완료되면 output_dir 하위에 best_model.pt 체크포인트 파일이 저장됩니다. torch.load()를 사용하여 model_state_dict를 읽어온 뒤 새로운 MNISTLogitNet 인스턴스에 로드하여 가중치를 재복원할 수 있습니다.

In [ ]:

import os

ckpt_path = "./result/mnist/best_model.pt"
if os.path.exists(ckpt_path):
    checkpoint = torch.load(ckpt_path, weights_only=True)
    
    # 이식 가능한 state_dict 기반 모델 가중치 복원
    eval_model = MNISTLogitNet()
    eval_model.load_state_dict(checkpoint["model_state_dict"])
    eval_model.eval()

    print("Loaded Checkpoint Version:", checkpoint.get("version"))
    print("Checkpoint Best Score:     ", checkpoint.get("score"))
    print("Execution Target Device:   ", checkpoint.get("device"))