[파이썬, Python] 딥러닝 - 3️⃣ 다중 분류 신경망 모델 구현하기!

📄 포켓몬 149종 분류 데이터셋

Train: https://www.kaggle.com/datasets/thedagger/pokemon-generation-one
Validation: https://www.kaggle.com/datasets/hlrhegemony/pokemon-image-dataset (디렉토리 비교 후 Train 데이터에 있는 클래스만 전처리함)

Complete Pokemon Image Dataset

 

Pokemon Generation One

 

---

1.  포켓몬 149종 분류

 

✅ 데이터셋 다운로드

import os
os.environ['KAGGLE_USERNAME'] = '****'  # username
os.environ['KAGGLE_KEY'] = '****'  # token key

!kaggle datasets download -d thedagger/pokemon-generation-one
!kaggle datasets download -d hlrhegemony/pokemon-image-dataset

# 압축 해제
!unzip -q pokemon-generation-one.zip
!unzip -q pokemon-image-dataset.zip

 

✅ 디렉토리 이름 변경하기 - 리눅스 문법으로

#디렉토리 이름 변경 - 리눅스 문법
!mv dataset train
!rm -rf train/dataset
!mv images validation

 

✅ 디렉토리 이름 변경하기 - 파이썬 shutil 

•  train: 149 classes, validation: 898 classes
• validation에서 train에 있는 디렉토리를 확인하여 없는 디렉토리는 제거(또는 이동)

import shutil

val_path = 'validation'
train_path = 'train'

val_dir = os.listdir(val_path)
train_dir = os.listdir(train_path)
# val_dir
for directory in val_dir:
  if directory not in train_dir:
    directory_path = os.path.join(val_path, directory)  # validation/폴더명
    shutil.rmtree(directory_path)

# 제거 후 폴더명 다시 갱신
train_dir = os.listdir(train_path)
val_dir = os.listdir(val_path)

print(len(val_dir), len(train_dir))

>>> 149 149

 

 ✅validation에 없는 2종의 포켓몬 디렉토리를 추가로 만들고 인터넷에서 2종의 포켓몬 이미지를 최소 2개 이상 넣어줌

folderlist = []
# validation에 없는 2종 찾아내기
for folder in train_dir:
  if folder not in val_dir:
    print(folder)
    folderlist.append(folder)

# validation에 해당 폴더 만들기
for folder in folderlist:
  print(folder)
  directory_path = os.path.join(val_path, folder)
  os.makedirs(directory_path, exist_ok=True)
# 추가 후 폴더명 다시 갱신
train_dir = os.listdir(train_path)
val_dir = os.listdir(val_path)

print(len(val_dir), len(train_dir))
>>> 149 149

 

✅ Farfetchd와 MrMimi 이미지를 구글에 검색하여 넣기 (2개씩 넣음)

▲ Farfetchd

▲ MrMime

 

✅ 패키지 로드

import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
from torchvision import datasets, models, transforms
from torch.utils.data import DataLoader
from torch.nn import functional as F

 

✅ GPU 사용하기

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

>>> cuda

 

✅ transforms 만들기

• Compose를 사용하여 사이즈, Affine, RandomHorizontalFlip, ToTensor 역할

data_transforms = {
    'train': transforms.Compose([
        transforms.Resize([224, 224]),
        transforms.RandomAffine(0, shear=10, scale=(0.8 , 1.2)),  # 랜덤하게 변경할 것을 선택(인덱스 0번부터 10가지 선택, 크기는 범위 +-20%하여 랜덤하게 변경)
        transforms.RandomHorizontalFlip(),  # 랜덤하게 이미지 좌우 반전
        transforms.ToTensor()  # 이미지를 텐서형으로 변환
    ]),
    'validation': transforms.Compose([
        transforms.Resize([224, 224]),  # 사이즈 맞춤
        transforms.ToTensor()
    ])
}

 

✅ 데이터셋 만들기

image_datasets = {
    'train': datasets.ImageFolder('train', data_transforms['train']), # data 폴더 안에 train 폴더를 데이터셋화
    'validation': datasets.ImageFolder('validation', data_transforms['validation'])
}

 

✅ 데이터로더 만들기

dataloaders ={
    'train': DataLoader(
        image_datasets['train'],
        batch_size=32,
        shuffle=True
  ),
    'validation':DataLoader(
        image_datasets['validation'],
        batch_size=32,
        shuffle=False
    )
}

print(len(image_datasets['train']), len(image_datasets['validation']))
>>> 10657 663

 

✅ 1개의 batch만큼 이미지를 출력

imgs, labels = next(iter(dataloaders['train']))


fig, axes = plt.subplots(4, 8, figsize=(20, 10))

for img, label, ax in zip(imgs, labels, axes.flatten()):
  ax.set_title(label.item())
  ax.imshow(img.permute(1,2,0))  # 텐서에 저장되어있을 때 shape(컬러, 가로, 세로) -> matplotlib에서는 (가로, 세로, 컬러채널)
  ax.axis('off')

▲ 배치사이즈 32만큼 32개의 이미지가 transform되어 나옴

 

 

✅ 사전 학습된 EfficientNetB4 모델 사용(IMAGENET1K_V1로 사전학)

model = models.efficientnet_b4(weights='IMAGENET1K_V1').to(device)
print(model)

▲ 모델의 정보 중 일부(FC레이어 부분을 확인!)

 

✅ fc 모델만 수정하기(ouput이 149종)

for param in model.parameters():
  param.requires_grad = False  # 가져온 파라미터(W, b)를 업데이트 하지 않음

# 모델의 FC 레이어이름이 classifier
model.classifier = nn.Sequential(
    nn.Linear(1792,512),
    nn.ReLU(),
    nn.Linear(512,149)  # output이 149개
).to(device)

📍 모델 수정부분에서 layer을 좀 더 쌓으면 Accuracy 증가시킬 수 있다.

 

 

✅ 모델 학습

# 학습
optimizer = optim.Adam(model.parameters(), lr=0.001)

epochs = 5

for epoch in range(epochs+1):
  for phase in ['train', 'validation']:  # train과  validation 따로 반복문을 돌아
    if phase == 'train':
      model.train()
    else:
      model.eval()   # 학습 모드에 있던 메모리를 지우고 바로 Test모드(훨씬 빠름)

    sum_losses = 0
    sum_accs = 0

    for x_batch, y_batch in dataloaders[phase]:  # train이라면 train에 대한 데이터로더, validataion이라면 validation에 대한 데이터로더 (따로 쓰지 않고 합쳐서 씀)
      x_batch = x_batch.to(device)
      y_batch = y_batch.to(device)

      y_pred = model(x_batch)

      loss = nn.CrossEntropyLoss()(y_pred, y_batch.long())

      if phase == 'train':
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

      sum_losses = sum_losses + loss.item()

      y_prob = nn.Softmax(1)(y_pred)
      y_pred_index = torch.argmax(y_prob, axis=1)
      acc = (y_batch == y_pred_index).float().sum() / len(y_batch) * 100
      sum_accs = sum_accs + acc.item()

    avg_loss = sum_losses / len(dataloaders[phase])
    avg_acc = sum_accs / len(dataloaders[phase])

    print(f'{phase:10s}: Epoch {epoch+1:4d}/{epochs}, Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.2f}%')

▲ epoch수를 늘려서 학습하면 Accuracy 상승, 시간문제로 5 epochs 돌

✅ 학습된 모델 파일 저장

torch.save(model.state_dict(), 'model.h7')

 

✅ 저장된 모델 학습

model = models.efficientnet_b4().to(device)

model.classifier = nn.Sequential(
    nn.Linear(1792,512),
    nn.ReLU(),
    nn.Linear(512,149)  # output이 149개
).to(device)

model.load_state_dict(torch.load('model.h7'))
model.eval()  # 평가모드로 사용

 

 

✅ 테스트(validation에 있는 2종의 포켓몬을 통해 분류테스트)

from PIL import Image  # 이미지를 코랩에서 띄우기


img1 = Image.open('/content/validation/Butterfree/0.jpg')   # validation에 있는 이미지 쓰기
img2 = Image.open('/content/validation/Charmeleon/2.jpg')

fig, axes = plt.subplots(1,2, figsize=(12,6))
axes[0].imshow(img1)
axes[0].axis('off')
axes[1].imshow(img2)
axes[1].axis('off')
plt.show()

 

✅ validation 이미지를 transform(텐서형으로 변환)

img1_input = data_transforms['validation'](img1)  # data_transforms의 validation 키의 value 값에 img1을 통과시킴 -> Resize되고 Tensor로 변환
img2_input = data_transforms['validation'](img2)
print(img1_input.shape)
print(img2_input.shape)

>>> torch.Size([3, 224, 224])
torch.Size([3, 224, 224])

 

✅ 두 이미지를 첫번째 차원에 맞춰 하나의 이미지 텐서로 결합

test_batch = torch.stack([img1_input, img2_input])  
test_batch = test_batch.to(device)
test_batch.shape   # 두 개 이미지가 붙음, torch.Size([2, 3, 224, 224]): (2(배치크기), 컬러채널, 세로, 가로)
>>> torch.Size([2, 3, 224, 224])

✅ 예측값 도출

y_pred = model(test_batch)
y_pred

 ✅예측값에 대한 예측 확률

y_prob = nn.Softmax(1)(y_pred)
y_prob

 

 

✅ 확률이 높은 k개의 데이터 뽑기

• prob이나 가중치가 많을 때 k개만 위에서부터 뽑음 ->  indices에는 인덱스, probs에는 값이 반환

probs, indices = torch.topk(y_prob, k=3, axis=-1)  # 상위 3개
probs = probs.cpu().data.numpy()  # cpu 계산
indices  =indices.cpu().data.numpy()  # cpu 계산
print(probs)
print(indices)

>>> [[9.9982733e-01 8.1344311e-05 2.8760607e-05]
 [5.5579984e-01 2.2686876e-01 8.2384981e-02]]
[[ 10 133 132]
 [ 15  74  14]]

 

✅ 예측 결과 시각화

fig, axes = plt.subplots(1,2, figsize=(12,6))

axes[0].set_title('{:.2f}% {},{:.2f}% {},{:.2f}% {}'.format(
    probs[0, 0] * 100, image_datasets['validation'].classes[indices[0,0]],
    probs[0, 1] * 100, image_datasets['validation'].classes[indices[0,1]],
    probs[0, 2] * 100, image_datasets['validation'].classes[indices[0,2]],
))
axes[0].imshow(img1)
axes[0].axis('off')

axes[1].set_title('{:.2f}% {},{:.2f}% {},{:.2f}% {}'.format(
    probs[1, 0] * 100, image_datasets['validation'].classes[indices[1,0]],
    probs[1, 1] * 100, image_datasets['validation'].classes[indices[1,1]],
    probs[1, 2] * 100, image_datasets['validation'].classes[indices[1,2]],
))
axes[1].imshow(img2)
axes[1].axis('off')
plt.show()

▲ 예측 확률과 포켓몬 이름