mnist.py

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from torch.utils.data import random_split
import matplotlib.pyplot as plt
import numpy as np

import os
os.environ["KMP_DUPLICATE_LIT_OK"]="TRUE"

# 오토인코더 모델 클래스 생성
class Autoencoder(nn.Module):
	def __init__(self):
		super(Autoencoder, self).__init__()
		
		#Encoder
		self.encoder = nn.Sequential(
			nn.Linear(28*28, 128),
			nn.ReLU(),
			nn.Linear(128, 64),
			nn.ReLU(),
			nn.Linear(64, 16),
			nn.ReLU(),
			nn.Linear(16, 2)	# 잠재공간 latent space
		)
		
		#Decoder
		self.decoder = nn.Sequential(
			nn.Linear(2, 16),
			nn.ReLU(),
			nn.Linear(16, 64),
			nn.ReLU(),
			nn.Linear(64, 128),
			nn.ReLU(),
			nn.Linear(128, 28*28),
			nn.Sigmoid()
		)
	# 순전파
	def forward(self, x):
		x_encoded = self.encoder(x) # x_encoded: 입력으로 주어진 784차원을 잠재 차원으로 줄인 값
		x = self.decoder(x_encoded)
		return x, x_encoded

# =============================
# --- MNIST 데이터셋 학습 준비 ---
# =============================
# 데이터 변환
transform = transforms.Compose([
	transforms.ToTensor(),
])

# 원본 학습 데이터셋 로드
full_train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)

# 학습 데이터셋을 train, validate로 분리 (8:2)
train_size = int(0.8 * len(full_train_dataset))
val_size = len(full_train_dataset) - train_size
train_subset, val_subset = random_split(full_train_dataset, [train_size, val_size])

# 학습데이터, 검증데이터 로드 
train_loader = DataLoader(train_subset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_subset, batch_size=1000, shuffle=False)

# 테스트 데이터 로드
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=True)

# ============================
# --- AutoEncoder 학습 진행 ---
# ============================
print(f"\n===== Training Autoencoder with 2D Latent Space =====")

# 모델, 손실 함수, 옵티마이저 초기화
model = Autoencoder()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
num_epochs = 60

train_losses = []
val_losses = []

# 학습 루프
for epoch in range(num_epochs):
    model.train()
    running_train_loss = 0.0
    for data in train_loader:
        img, _ = data
        img = img.view(img.size(0), -1)
        output, _ = model(img)
        loss = criterion(output, img)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        running_train_loss += loss.item()
    
    epoch_train_loss = running_train_loss / len(train_loader)
    train_losses.append(epoch_train_loss)

    # --- 검증 단계 ---
    model.eval()
    running_val_loss = 0.0
    with torch.no_grad():
        for data in val_loader:
            img, _ = data
            img = img.view(img.size(0), -1)
            output, _ = model(img)
            loss = criterion(output, img)
            running_val_loss += loss.item()
    
    epoch_val_loss = running_val_loss / len(val_loader)
    val_losses.append(epoch_val_loss)

    # 에포크마다 손실과 검증 수치 출력
    print(f'Epoch [{epoch+1}/{num_epochs}], Train Loss: {epoch_train_loss:.4f}, Val Loss: {epoch_val_loss:.4f}')

print("Training Complete")

# 손실 시각화
plt.figure(figsize=(10, 5))
plt.plot(train_losses, label='Training Loss')
plt.plot(val_losses, label='Validation Loss')
plt.title('Training & Validation Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.grid(True)
plt.show()

# ==========================================================
# --- 모델 복원도 평가 (테스트 데이터셋 전체에 대한 평균 손실 계산) ---
# ==========================================================
model.eval()
test_loss = 0.0
with torch.no_grad():
    for data in test_loader:
        imgs, _ = data
        imgs = imgs.view(imgs.size(0), -1)
        outputs, _ = model(imgs)
        loss = criterion(outputs, imgs)
        test_loss += loss.item()
test_loss /= len(test_loader)
print(f"\n--- Results for 2D Model ---")
print(f"복원도 (Average Test Loss): {test_loss:.4f}\n")


print(f"Visualizing reconstructed images for 2D model...")

# ================================
# --- 실제 데이터셋을 이용한 테스트 ---
# ================================
# 테스트 데이터
dataiter = iter(test_loader)
images, _ = next(dataiter)
images_flattened = images.view(images.size(0), -1)

# 테스트 데이터를 학습시킨 AutoEncoder를 통해 테스트
output, _ = model(images_flattened)
output = output.view(output.size(0), 1, 28, 28).detach()

# 결과 시각화
fig, axes = plt.subplots(nrows=2, ncols=10, sharex=True, sharey=True, figsize=(20,4))
fig.suptitle(f'Reconstruction Results (2D Latent Space)')
# 원본 이미지
for i in range(10):
    img = images[i].squeeze()
    axes[0,i].imshow(img.numpy(), cmap='gray')
    axes[0,i].get_xaxis().set_visible(False)
    axes[0,i].get_yaxis().set_visible(False)
    if i == 0: axes[0,i].set_title('Originals', loc='left')

# 복원된 이미지
for i in range(10):
    img = output[i].squeeze()
    axes[1,i].imshow(img.numpy(), cmap='gray')
    axes[1,i].get_xaxis().set_visible(False)
    axes[1,i].get_yaxis().set_visible(False)
    if i == 0: axes[1,i].set_title('Reconstructed', loc='left')
plt.show()

# ============================
# --- latent vactors 시각화 ---
# ============================
model.eval()	# 모델을 평가 모드로 변경

latent_vectors = []
labels = []

with torch.no_grad():
	for data in test_loader:
		imgs, lbls = data
		imgs = imgs.view(imgs.size(0), -1)
		_, latents = model(imgs)
		latent_vectors.append(latents)
		labels.append(lbls)
		
latent_vectors = torch.cat(latent_vectors, dim=0)
labels = torch.cat(labels, dim=0)

# 2D 시각화
plt.figure(figsize=(12, 10))
# 각 숫자 레이블(0-9)에 대해 다른 색상으로 산점도 그리기
for i in range(10):
	indices = labels == i
	plt.scatter(latent_vectors[indices, 0], latent_vectors[indices, 1], label=str(i), alpha=0.7, s=15)

plt.xlabel('Latent X')
plt.ylabel('Latent Y')
plt.title('2D Latent Space Visualization')
plt.legend()
plt.grid(True)
plt.show()

# =======================================
# --- 3차원 이상의 latent vector 확인 시 ---
# =======================================
"""
from sklearn.manifold import TSNE

# t-SNE 적용
tsne = TSNE(n_components=2, perplexity=30, max_iter=300)
tsne_result = tsne.fit_transform(latent_vectors)

# 시각화
plt.figure(figsize=(20, 8))

# t-SNE 결과 플롯
for i in range(10):
    indices = labels == i
    plt.scatter(tsne_result[indices, 0], tsne_result[indices, 1], label=str(i), alpha=0.7, s=15)
plt.title('t-SNE Visualization of Latent Space')
plt.xlabel('t-SNE Dimension 1')
plt.ylabel('t-SNE Dimension 2')
plt.legend()
plt.grid(True)

plt.show()
"""