[파이토치] CNN 모델 학습, 평가, 추론

https://github.com/stonegyoung/Pytorch_Study/blob/master/lenet_%EA%B5%AC%ED%98%84%2BMNIST.ipynb

Pytorch_Study/lenet_구현+MNIST.ipynb at master · stonegyoung/Pytorch_Study

---

 

이미지 준비

전처리까지 완료된 train_loader와 test_loader를 가지고 진행합니다.

2024.12.31 - [Pytorch] - [파이토치] MNIST로 이미지 데이터 처리

[파이토치] MNIST로 이미지 데이터 처리

 

모델 만들기

lenet5 아키텍처로 진행합니다.

2025.01.04 - [Pytorch] - [파이토치] CNN Lenet5 구현

[파이토치] CNN Lenet5 구현

---

이제 MNIST 데이터를 학습 시켜보겠습니다.

 

모델 학습을 위한 최적화함수, 손실함수, 에폭 정의

from torch.optim import Adam
import torch.nn as nn

# 최적화함수
optimizer = Adam(model.parameters(), lr=0.001)
# 손실함수
criterion = nn.CrossEntropyLoss()
# 에폭
epochs = 10

 

gpu 확인

device = 'cuda' if torch.cuda.is_available() else 'cpu'
device

 

저는 'cuda'가 떴기 때문에 gpu를 사용할 수 있습니다.

그럼 model을 gpu로 옮겨주겠습니다.

model.to(device)

 

---

모델 학습

train_loader로 매 배치마다

기울기 0으로 초기화 -> 순전파 -> 손실함수 계산 -> 역전파 -> 최적화함수로 모델 파라미터 업데이트

 

이걸 정해둔 에폭만큼 반복하면 됩니다

 

+ 배치마다 기울기를 0으로 초기화하는 이유
-> 각 배치마다 독립적으로 파라미터 업데이트가 이루어져야 하기 때문.
즉, 각 배치에서 계산된 기울기를 그 배치에만 적용하고, 다음 배치의 기울기가 이전 배치의 영향을 받지 않도록 해야 함

 

model.train()

loss_list = []
min_loss = int(1e9)
for epoch in range(epochs):
  losses = 0
  for data, target in train_loader:
    optimizer.zero_grad() # 배치마다 기울기 초기화(각 배치의 정보에 맞게 가중치 업데이트를 위해)
    y_pred = model(data.to(device))
    loss = criterion(y_pred, target.to(device))
    losses += loss.item()
    loss.backward()
    optimizer.step()
  avg_loss = losses/len(train_loader) # 배치 별 손실 값 평균
  loss_list.append(avg_loss) # 에폭 별 손실 값 리스트에 저장
  print(f'epoch: {epoch}, loss: {avg_loss}')
  if avg_loss < min_loss: # 손실 값이 최저일 때만 저장
    torch.save(model.state_dict(), f'model{epoch}.pth') # 모델 가중치 저장
    min_loss = avg_loss # 최소 손실 값 갱신

 

시각화

import matplotlib.pyplot as plt

plt.plot(range(epochs), loss_list)
plt.show()

 

---

모델 검증

model.eval(): 모델 평가 모드로 전환

with torch.no_grad(): 기울기 계산되지 않도록 방지

 

test_loader로 매 배치마다

출력값(각 클래스에 대한 예측값) 중 제일 값이 높은 argmax값 구하기 -> 정답과 비교

 

model.eval() # 모델 평가 모드로 전환
correct = 0 # 맞춘 수
total = 0 # 데이터 수

with torch.no_grad(): # 기울기 계산 방지
  for data, target in test_loader:
    y_pred = model(data.to(device))
    result = y_pred.max(dim=1)[1] # 확률이 가장 높은 클래스 선택
    correct += sum(result==target.to(device)).item() # 맞춘 수
    total += target.size(0)
print(f'accuracy: {correct/total}') # 정확도 출력

 

---

모델 추론

이제 제가 쓴 손글씨 이미지를 가지고 추론해보겠습니다.

이미지 확인

이미지를 PIL.Image로 가져오겠습니다.

# 이미지 불러오기
from PIL import Image
import numpy as np

img = Image.open('mnist_infer.png')
img

 

numpy로 바꿔 shape을 확인해보겠습니다.

numpy_img = np.array(img)
numpy_img.shape

png이기 때문에 채널 수가 4(RGBA)

jpg는 채널 수가 3(RGB)

 

또한 컬러 -> 흑백(1 채널)로 바꿔줘야 합니다.

 

전처리

모델에 넣기 위해 모델에 맞는 이미지로 전처리 해줘야 합니다.

- 알파 채널 제거 (png일 경우)

transforms.Lambda(lambda x: x[:3])

- 컬러 -> 흑백으로 변경

transforms.Grayscale(num_output_chnnels=1)

from torchvision import transforms

# 전처리
data_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((32,32)),
    transforms.Lambda(lambda x: x[:3]), # rgba -> rgb
    transforms.Grayscale(num_output_channels=1), # 3채널 -> 1채널(흑백으로)
    transforms.Normalize((0.5,),(1.0,))
])

 

전처리가 적용된 이미지를 transform_img라는 변수로 두겠습니다.

transform_img = data_transform(img)

 

전처리 적용 추론 이미지 shape 확인

transform_img.shape

 

이제 전처리 적용 추론 이미지 시각화를 해보겠습니다.

시각화 함수는 그동안 한 것과 동일합니다

import matplotlib.pyplot as plt

def imgshow(tensor_array, cmap=None):
  # tensor -> numpy
  numpy_array = tensor_array.numpy()
  # 차원 변환 (channel, h, w) -> (h, w, channel)
  num_img = numpy_array.transpose(1,2,0) # (0, 1, 2) -> (1, 2, 0)
  plt.axis('off') # 눈금 제거
  plt.imshow(num_img, cmap=cmap)

imgshow(transform_img, 'gray')

model은 cpu로 내리겠습니다.

model.to('cpu')

 

 

모델 추론 함수

배치 사이즈도 없기 때문에 배치 사이즈 차원을 추가해줘야 합니다.

모델 추론도 검증과 동일하게 모델 평가 모드와 기울기 계산 방지를 해야 합니다.

def inference(img):
  img = img.unsqueeze(dim=0) # 맨 앞에 배치 사이즈 추가
  model.eval()

  with torch.no_grad():
    pred = model(img)
    result = pred.argmax(dim=1)
    print(f'predict: {result.item()}')

 

inference(transform_img)