[Model] MNIST 문자 인식 모델 구조화 및 설명

모델 학습에 있어 각 기능별로 클래스화시켜 설명할 예정이다.

일단 두 가지로 파일을 나누어 모델을 학습시키는 파일(mnist_model.py)과,

모델을 평가하고 테스트하는 파일(mnist_model_eval.py)로 나누었다.

 

 

 

개발 환경

- os : window

- virtual env : anaconda

- python 3.10.16

mnist_model.py

0. 사용 라이브러리

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

import os
import torch.onnx

1. GPU_manager : GPU 설정

- (굳이 만들필요는 없었던 것 같지만)

- 기본 device를 CPU로 설정하였고,

- setDevice 멤버 함수로 cuda가 가능하면 device를 cuda로 설정

class GPUManager:
    def __init__(self):
        self.device_ = 'cpu'
        torch.manual_seed(777)

    def setDevice(self):
        self.device_ = 'cuda' if torch.cuda.is_available() else 'cpu'
        if self.device_ == 'cuda':
            torch.cuda.manual_seed_all(777)
            print(f"{self.device_} is available")
        return self.device_

2. HyperParameters

class HyperParameters:
    def __init__(self):
        self.learning_rate_ = 0.001 	# 모델 학습시 파라미터의 변화율
        self.batch_size_ = 64		# 한 번에 처리할 이미지의 수
        self.num_classes_ = 10		# MNIST 데이터셋 출력층 크기(분류 10가지)
        self.epochs_ = 10		# 전체 데이터셋을 몇 번 반복학습

3. MNISTDataloader

transforms.ToTensor()

- 입력이미지를 tensor로 변환(픽셀값을 0~1로 정규화)

- 차원 순서를 (H,W,C)에서 (C,H,W)로 변경

- ouput: (1,28,28)

 

DataLoader()

- 이미지 텐서를 학습시 epoch에서 얼만큼 처리될 지 결정

- output : image tensor(batch_size, channel, height, width)

- ex) 전체 데이터수 60000 , batch_size 64 => 938번 배치가 진행

 

class MNISTDataLoader:
    def __init__(self, hp: HyperParameters):
        self.hp_ = hp
        self.transform_ = transforms.Compose([transforms.ToTensor()]) # 입력이미지 tensor로 변환
        self.data_path_ = './data/MNIST'

    # MNIST train/test 데이터셋 로드
    def loadDatasets(self):
        self.train_set_ = torchvision.datasets.MNIST(
            root=self.data_path_,
            train=True,
            download=True,
            transform=self.transform_
        )

        self.test_set_ = torchvision.datasets.MNIST(
            root=self.data_path_,
            train=False,
            download=True,
            transform=self.transform_
        )
        return self

    def createDataLoaders(self):
        if self.train_set_ is None or self.test_set_ is None:
            raise ValueError("데이터셋이 로드되지 않았습니다. loadDatasets()를 먼저 호출하세요.")

        # image tensor(batch_size, channel, height, width)
        self.train_loader_ = DataLoader(
            dataset=self.train_set_,
            batch_size=self.hp_.batch_size_,
            shuffle=True
        )

        self.test_loader_ = DataLoader(
            dataset=self.test_set_,
            batch_size=self.hp_.batch_size_,
            shuffle=False
        )
        return self

 

 

여러 데이터 증강 옵션

transforms.RandomRotation(degrees=10) # 회전

transforms.RandomAffine(degrees=10, translate=(0.1, 0.1)) # 아핀 변환

transforms.RandomHorizontalFlip(p=0.5) # 좌우 반전

transforms.RandomCrop(28, padding=4) # 패딩 후 크롭

 

4. ConvNet

 

• 입력층
  • 입력 크기: (64, 1, 28, 28)
  • 배치 크기: 64
  • 채널: 1 (흑백 이미지)
  • 이미지 크기: 28x28
• 첫 번째 컨볼루션 블록 (self.conv1_)
  • Conv2d: (64,1,28,28) → (64,10,24,24)
  • ReLU: 활성화 함수
  • MaxPool2d: (64,10,24,24) → (64,10,12,12)
• 두 번째 컨볼루션 블록 (self.conv2_)
  • Conv2d: (64,10,12,12) → (64,20,8,8)
  • ReLU: 활성화 함수
  • MaxPool2d: (64,20,8,8) → (64,20,4,4)
• Dropout2D
  • 크기 유지: (64,20,4,4)
  • 25% 확률로 특징 맵 전체를 드롭아웃
  • 과적합 방지
• Flatten
  • 변환: (64,20,4,4) → (64,320)
  • 완전연결층을 위한 1차원 벡터화
• 완전연결층 (self.fc_layers_)
  • 첫 번째 FC: (64,320) → (64,100)
  • ReLU 활성화
  • 두 번째 FC: (64,100) → (64,10)
• 출력층
  • Log_Softmax: (64,10) 유지
  • 각 숫자(0-9)에 대한 확률 분포 출력

 

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        # Conv1 -> ReLU -> Max Pool
        self.conv1_ = nn.Sequential(
            nn.Conv2d(1, 10, kernel_size=5),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )

        # Conv2 -> ReLU -> Max Pool
        self.conv2_ = nn.Sequential(
            nn.Conv2d(10, 20, kernel_size=5),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )

        # dropout, 25% 확률로 무작위 뉴런 비활성화           
        self.dropout_ = nn.Dropout2d(p=0.25)           

        self.fc_layers_ = nn.Sequential(
            nn.Linear(320, 100),
            nn.ReLU(),
            nn.Linear(100, 10)
        )

    def forward(self, x):
        # 컨볼루션 레이어 통과
        x = self.conv1_(x)
        x = self.conv2_(x)

        # 드롭아웃 적용
        x = self.dropout_(x)

        # Flatten
        x = x.view(x.size(0), -1)

        # Fully Connected layer 
        x = self.fc_layers_(x)

        # Log Softmax 적용
        x = F.log_softmax(x, dim=1)

        return x

 

5. Learning

 

trainStep

- 퍼셉트론과 같은 가중치 조정 과정을 보여준다.

- output : 한 퍼셉트론의 loss

trainEpoch

- 한 epoch에서 여러 trainStep이 있으므로 개념을 동일하게 사용

- 그리고 한 epoch라는 거시적 관점에서 Step의 손실 평균을 낸다.

- output : 한 epoch 동안의 avg_loss

 

train

- 여러 epoch에서 loss가 줄어들고 있다면 그것은 올바른 방향으로 학습 중임을 나타낸다.

class Learning:
    def __init__(self, model, optimizer, criterion, device, epochs):
        self.model_ = model
        self.optimizer_ = optimizer
        self.criterion_ = criterion
        self.device_ = device
        self.epochs_ = epochs

        self.train_loss_ = []
        self.test_loss_ = []
        self.is_trained_ = False

    def trainStep(self, data, target):
        data, target = data.to(self.device_), target.to(self.device_)
        self.optimizer_.zero_grad()
        
        hypothesis = self.model_(data)			# 순전파
        cost = self.criterion_(hypothesis, target) 	# 손실 계산
        cost.backward() 				# 역전파 계산
        self.optimizer_.step()				# 파라미터 업데이트

        return cost

    def trainEpoch(self, train_loader):
    
        self.model_.train()  # 학습 모드 설정
        avg_cost = 0

        for data, target in train_loader:
            cost = self.trainStep(data, target)
            avg_cost += cost.item() / len(train_loader)

        self.train_loss_.append(avg_cost)
        return avg_cost

    def train(self, train_loader):

        for epoch in range(self.epochs_):
            avg_cost = self.trainEpoch(train_loader)
            print('[Epoch: {:>4}] cost = {:>.9}'.format(epoch + 1, avg_cost))

        self.is_trained_ = True
        return self.train_loss_

 

6. ModelSaver

(made by AI)

class ModelSaver:
    def __init__(self, model_dir='saved_models'):
        self.model_dir = model_dir
        os.makedirs(model_dir, exist_ok=True)

    def save_model(self, model, filename):
        """모델 저장"""
        save_path = os.path.join(self.model_dir, filename)
        torch.save(model.state_dict(), save_path)
        print(f'Model saved to {save_path}')

    def load_model(self, model, filename):
        """모델 불러오기"""
        load_path = os.path.join(self.model_dir, filename)
        model.load_state_dict(torch.load(load_path))
        print(f'Model loaded from {load_path}')
        return model

7. Main

실제로는 하이퍼파라미터를 조정하여 진행하였고, model을 저장하였다.

if __name__ == '__main__':
    # GPU 사용 여부
    gpu_manager = GPUManager()
    device = gpu_manager.setDevice()
    
    
    # 하이퍼파라미터 설정
    hp = HyperParameters(batch_size= 256, epoch = 15)
    
    # 데이터로더 생성
    data_loader = MNISTDataLoader(hp)
    data_loader.loadDatasets().createDataLoaders()
    
    # 데이터로더 가져오기
    train_loader, test_loader = data_loader.getLoaders()
    
    # 모델 설정 및 GPU로 이동
    model = ConvNet().to(device)

    # Learning 클래스 인스턴스 생성
    learning = Learning(model = model, 
                        optimizer= optim.Adam(model.parameters(), lr=hp.learning_rate_), 
                        criterion= nn.CrossEntropyLoss(),
                        device = device)
    
    # 학습 실행
    train_losses = learning.train(train_loader)
    
    
    # 1. PyTorch 모델 저장
    model_saver = ModelSaver()
    model_saver.save_model(learning.model_, 'mnist_trained_model.pth')
    
    # 2. ONNX 형식으로 내보내기
    export_to_onnx(learning.model_, 'mnist_model.onnx')

 

 

 

---

 

mnist_model_eval.py

import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

# 기존 클래스들 import

from mnist_model import ConvNet, HyperParameters, MNISTDataLoader

 

1. ModelEvaluator

loadModel

- 일단 CNN 구조를 외부 클래스에서 가져오기

- 저장된 가중치 가져오기

- 평가 모드로 설정 (dropout, batch normalization을 비활성화)

 

ModelEvaluator

- 모델 예측 후, 가장 높은 확률을 가지는 클래스 선택

 

class ModelEvaluator:
    def __init__(self, model_path, device):
        self.model_path_ = model_path
        self.device_ = device
        self.model_ = None
        
    def loadModel(self):
        """저장된 모델 불러오기"""
        self.model_ = ConvNet().to(self.device_)
        self.model_.load_state_dict(torch.load(self.model_path_, weights_only=True))
        self.model_.eval()  # 평가 모드로 설정
        print(f"Model loaded from {self.model_path_}")
        
    def evaluate(self, test_loader):
        """모델 평가"""
        if self.model_ is None:
            raise ValueError("먼저 모델을 로드해주세요.")
            
        correct = 0
        total = 0
        
        with torch.no_grad():  # 기울기 계산 비활성화
            for data, target in test_loader:
                data, target = data.to(self.device_), target.to(self.device_)
                
                outputs = self.model_(data)
                _, predicted = torch.max(outputs.data, 1)
                
                total += target.size(0)
                correct += (predicted == target).sum().item()
                
        accuracy = 100 * correct / total
        print(f'Test Accuracy: {accuracy:.2f}%')
        return accuracy

 

 

 

2. ModelRealTester:

(made by AI)

class ModelRealTester:
   def __init__(self, model, device):
       self.model_ = model
       self.device_ = device
       self.fig_size_ = (12, 6)
       self.num_images_ = 10

   def visualizePredictions(self, test_loader):
       """모델 예측 결과 시각화"""
       # 모델을 평가 모드로 설정
       self.model_.eval()

       # 테스트 데이터 가져오기
       data_iter = iter(test_loader)
       images, labels = next(data_iter)
       images, labels = images.to(self.device_), labels.to(self.device_)

       # 예측 수행
       with torch.no_grad():
           outputs = self.model_(images)
           _, predicted = torch.max(outputs.data, 1)

       # 시각화
       fig = plt.figure(figsize=self.fig_size_)
       for idx in range(self.num_images_):
           ax = fig.add_subplot(2, 5, idx + 1)
           # CPU로 이동하고 채널 차원 제거
           ax.imshow(images[idx].cpu().squeeze(), cmap='gray')
           # 실제 레이블과 예측값 표시
           ax.set_title(f'True: {labels[idx].item()}\nPred: {predicted[idx].item()}')
           ax.axis('off')

       plt.tight_layout()
       plt.show()

   def setFigureSize(self, width, height):
       """Figure 크기 설정"""
       self.fig_size_ = (width, height)
       return self

   def setNumImages(self, num):
       """시각화할 이미지 수 설정"""
       self.num_images_ = num
       return self

   def saveFigure(self, test_loader, filename):
       """예측 결과를 이미지 파일로 저장"""
       self.model_.eval()
       data_iter = iter(test_loader)
       images, labels = next(data_iter)
       images, labels = images.to(self.device_), labels.to(self.device_)

       with torch.no_grad():
           outputs = self.model_(images)
           _, predicted = torch.max(outputs.data, 1)

       fig = plt.figure(figsize=self.fig_size_)
       for idx in range(self.num_images_):
           ax = fig.add_subplot(2, 5, idx + 1)
           ax.imshow(images[idx].cpu().squeeze(), cmap='gray')
           ax.set_title(f'True: {labels[idx].item()}\nPred: {predicted[idx].item()}')
           ax.axis('off')

       plt.tight_layout()
       plt.savefig(filename)
       plt.close()
       print(f"Figure saved as {filename}")

 

3. Main

if __name__ == "__main__":
    # GPU 설정
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 하이퍼파라미터 및 데이터 로더 설정
    hp = HyperParameters()
    data_loader = MNISTDataLoader(hp)
    data_loader.loadDatasets().createDataLoaders()
    _, test_loader = data_loader.getLoaders()
    
    # 평가 실행
    evaluator = ModelEvaluator('saved_models/mnist_trained_model.pth', device)
    evaluator.loadModel()
    
    accuracy = evaluator.evaluate(test_loader)

    # 시각화 도구 설정
    visualizer = ModelRealTester(evaluator.model_, device)
    
    # 기본 시각화
    visualizer.visualizePredictions(test_loader)
    
    # 설정 변경 후 시각화
    visualizer.setFigureSize(15, 8).setNumImages(10).visualizePredictions(test_loader)
    
    # 이미지로 저장
    visualizer.saveFigure(test_loader, 'predictions.png')