[ResNet] 이미지 분류기 만들기
ResNet모델을 이용하여 알약을 분류하기 위한 인공지능 모델을 만들어보자
나는 딥러닝, 머신러닝에 대해 거의 무지하기 때문에, ResNet모델에 대해 깊은 이해보다 모델을 학습시키고 사용하는 방법에 대해 연구하였고 그 과정에서 알게된 사실들을 정리하기 위해 이 글을 작성한다.
ResNet 모델은 resnet18, 34, 50, 등으로 구성되어 있으며 뒤에 붙은 숫자는 레이어의 수를 의미한다.
앞서 말했듯 나는 머신러닝에 대해 무지하기 때문에 레이어에 대한 설명은 gpt의 설명을 인용하도록 하겠다.
예를 들어, 한 레이어는 입력 이미지의 특정한 부분에 주목할 수 있도록 도와줄 수 있습니다. 다른 레이어는 그림의 모서리나 선을 감지하는 역할을 할 수 있습니다. 각 레이어는 입력으로부터 특정한 특징을 추출하고, 그 다음 레이어로 그 정보를 전달합니다. 레이어 간의 연결은 데이터의 흐름을 나타냅니다. 네트워크의 깊이는 레이어의 수에 따라 결정됩니다. 더 많은 레이어를 추가하면 네트워크는 더 복잡한 패턴을 학습할 수 있게 됩니다.
유튜브와 검색등을 통해 얻은 정보로는, resnet 모델은 레이어 수가 늘어날수록 더 복잡한 패턴을 학습가능하지만, 학습 시킨 클래스의 수가 레이어수보다 적다면, 오히려 정확도가 낮아지는 현상이 발생한다고 한다.
초기 계획은 AI Hub에서 다운받은 몇만 종류의 알약을 학습 시키고자 하였으나, 몇만 종류의 알약 이미지 데이터를 모두 다운받아 학습 시키는것은 상당히 많은 시간이 소요될 뿐더러, 불필요한 기능이라고 느껴져 기초 데이터에서 제공해주는 1000종류의 알약 데이터를 학습시키기로 변경하였고. gpt를 통해 검색한 결과 resnet50모델로 학습시키는것이 가장 적합하다는 답을 얻을 수 있었다.

resnet모델의 사용법에 대해 전혀 알지 못하기 때문에
1000개의 데이터를 바로 학습시키기 이전에, 6개의 모델을 학습시켜보면서 resnet모델을 사용하는 방법과, 해당 모델이 프로젝트에 적합한지 여부를 판단하고자 한다. 6개를 학습시킬것이므로 resnet모델은 resnet18을 사용하였다.
모델의 학습 방법은 유튜버 '동빈나'님의 이미지 분류기 만들기 영상을 참고하였다.
학습에 필요한 준비물은 YOLO와 동일하게 train과 test로 분리된 이미지들이 필요하지만, 객체의 위치 정보를 담은 바운딩 박스를 이용하여 학습하는 YOLO와 달리 이미지를 이용하여 특징을 판별하기 때문에 바운딩 박스가 필요하지 않다.
학습 데이터를 준비하기 위해 AI Hub에서 다운받은 파일을 전처리 하는 과정이 필요하다.
내가 보유중인 데이터는 아래와 같이 클래스별로 분리된 폴더에, 해당 클래스에 해당하는 이미지 1296장과 각 이미지의 정보를 담고 있는 json파일 1296장으로, 각 폴더마다 2592개의 데이터 파일이 존재한다.(이미지 + json)


우리는 이 파일을 적절히 분리하여
70퍼센트를 객체를 학습시키기 위한 train데이터로 삼고, 나머지 30퍼센트를 학습을 검증하기 위한 데이터로 분리해야한다.
직접 분리하기엔 적지않은 시간이 소요되므로, 파이썬 코드를 작성하여 png파일을 분리하였다.
파이썬 코드의 예시를 들기 위해 3개의 폴더를 대상으로 데이터 분리 작업을 하였다.
import os
import random
import shutil
# 경로 설정
source_dir = r"C:\proj_pill" # 데이터 파일이 존재하는 폴더의 위치
train_dir = r"C:\proj_pill\train" # train폴더 생성이후 코드에서 이용하기 위해
test_dir = r"C:\proj_pill\test" # test폴더 생성이후 코드에서 이용하기 위해
# 폴더 내의 파일들을 리스트로 저장
folders = os.listdir(source_dir)
# train, test 폴더 생성
os.makedirs(train_dir, exist_ok=True)
os.makedirs(test_dir, exist_ok=True)
# 각 폴더에 대해 데이터 분리
for folder in folders:
folder_path = os.path.join(source_dir, folder)
if not os.path.isdir(folder_path):
continue
# png 파일 분리 => json파일은 무시
png_files = []
for file in os.listdir(folder_path):
if file.lower().endswith(".png"):
png_files.append(file)
# 파일 리스트를 무작위로 섞음
random.shuffle(png_files)
# 폴더 내의 하위 폴더 생성
train_folder_path = os.path.join(train_dir, folder)
test_folder_path = os.path.join(test_dir, folder)
os.makedirs(train_folder_path, exist_ok=True)
os.makedirs(test_folder_path, exist_ok=True)
# 학습용 데이터 개수 계산
train_count = int(len(png_files) * 0.7)
# 학습용 데이터 잘라내기 및 이동
for file in png_files[:train_count]:
src_path = os.path.join(folder_path, file)
dst_path = os.path.join(train_folder_path, file)
shutil.move(src_path, dst_path)
# 테스트용 데이터 잘라내기 및 이동
for file in png_files[train_count:]:
src_path = os.path.join(folder_path, file)
dst_path = os.path.join(test_folder_path, file)
shutil.move(src_path, dst_path)
print("데이터 분리가 완료되었습니다.")
위 코드를 실행하면, 데이터 파일이 존재하는 폴더의 모든 폴더 이름을 읽어와 저장하고, 하위 폴더로 train폴더와 test폴더를 만들어 각 클래스별로 70:30 비율로 이미지를 분리하여 저장한다.
실행 결과를 확인하면 아래와 같다.




1296 * 0.7 = 907이므로, 907장의 이미지 데이터가 train폴더의 Tylenol폴더에 들어간 것을 확인할 수 있다.(K-004378이 원래 폴더의 이름이나, 편의를 위해 알약의 이름으로 변경하였다.)
이제 데이터가 준비되었으니 구글 드라이브에 위 데이터들을 업로드한다.
코랩은 런타임 연결시간이 오래되면 업로드한 파일을 다시 이용할 수 없으므로, 차후에 다시 이미지 데이터를 이용하기 위해 구글 드라이브와 마운트한다.
from google.colab import drive
drive.mount('/content/drive')
결과를 이미지로 확인하기 위해 matplotlib을 import하고, 학습에 필요한 라이브러리들을 import한다.
import matplotlib.pyplot as plt # 맷플롯립 import 하기(데이터 시각화용)
import torch
import torch.nn as nn
import torch.optim as optim
import os
import torchvision
from torchvision import datasets, models, transforms
import numpy as np
import time
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # device 객체
위 과정까지 정상적으로 실행하였다면 코랩의 화면은 아래와 같아야한다.
왼쪽 폴더에서 구글 드라이브에 마운트 된것을 확인할 수 있다.

transforms_train과 test는 학습의 정확도를 높이기 위한 전처리 과정으로 이미지에 여러가지 변형을 가한다.
아래 코드를 실행하여 각 폴더의 클래스별 이미지들을 불러와 저장하고, 학습시킬 클래스들의 이름을 class_names에 기록한다.
class_names 리스트가 존재하지 않으면, 추후 학습이 제대로 되었는지 확인할 수 없다.
# 데이터셋을 불러올 때 사용할 변형(transformation) 객체 정의
transforms_train = transforms.Compose([
transforms.Resize((224, 224)),
transforms.RandomHorizontalFlip(), # 데이터 증진(augmentation)
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # 정규화(normalization)
])
transforms_test = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
data_dir = './drive/MyDrive/pill_data' # 데이터가 저장된 폴더 (train,test폴더의 상위 디렉토리)
train_datasets = datasets.ImageFolder(os.path.join(data_dir, 'train'), transforms_train)
test_datasets = datasets.ImageFolder(os.path.join(data_dir, 'test'), transforms_test)
train_dataloader = torch.utils.data.DataLoader(train_datasets, batch_size=4, shuffle=True, num_workers=4)
test_dataloader = torch.utils.data.DataLoader(test_datasets, batch_size=4, shuffle=True, num_workers=4)
print('학습 데이터셋 크기:', len(train_datasets))
print('테스트 데이터셋 크기:', len(test_datasets))
class_names = train_datasets.classes
print('클래스:', class_names)
구글 드라이브에 pill_data라는 폴더를 만들어 그곳에 사전에 준비한 데이터를 넣어줬다.

코드를 실행하면 위와 같이 총 학습 데이터셍의 크기와, 테스트 데이터 셋의 크기, 클래스 목록들이 제대로 불려진 것을 확인할 수 있다.
시각적으로 결과를 확인하기 위해 이미지를 매개변수로 받아 화면에 출력해주는 imshow함수를 작성하고, 무작위로 이미지를 불러와 제대로 동작하는지 확인한다.
def imshow(input, title):
# torch.Tensor를 numpy 객체로 변환
input = input.numpy().transpose((1, 2, 0))
# 이미지 정규화 해제하기
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
input = std * input + mean
input = np.clip(input, 0, 1)
# 이미지 출력
plt.imshow(input)
plt.title(title)
plt.show()
# 학습 데이터를 배치 단위로 불러오기
iterator = iter(train_dataloader)
# 현재 배치를 이용해 격자 형태의 이미지를 만들어 시각화
inputs, classes = next(iterator)
out = torchvision.utils.make_grid(inputs)
imshow(out, title=[class_names[x] for x in classes])
코드를 실행하면 아래와 같이 무작위로 이미지를 가져와 화면에 보여준다.

이제 학습을 진행해보자.
전이학습을 할것이므로, 사전에 만들어진 resnet18모델이 필요하다. 아래 코드를 통해 model에 resnet18모델을 불러온다.
학습시킬 모델의 수는 6개이므로, 출력 뉴런 수를 6개로 변경한다.
from torchvision.models.resnet import ResNet18_Weights
model = models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
num_features = model.fc.in_features
# 전이 학습(transfer learning): 모델의 출력 뉴런 수를 6개로 교체하여 마지막 레이어 다시 학습
model.fc = nn.Linear(num_features, 6)
model = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
아래와 같이 Downloading이 되면 성공이다.

이제 이 model에 아래 코드를 이용하여 학습을 진행한다.
epochs는 각 이미지당 학습시킬 횟수를 의미하며 이 경우 50*907*6 번의 학습이 진행된다.(앞서 설명했듯 전문가가 아니기 때문에 정확하진 않다.)
num_epochs = 50
model.train()
start_time = time.time()
# 전체 반복(epoch) 수 만큼 반복하며
for epoch in range(num_epochs):
running_loss = 0.
running_corrects = 0
# 배치 단위로 학습 데이터 불러오기
for inputs, labels in train_dataloader:
inputs = inputs.to(device)
labels = labels.to(device)
# 모델에 입력(forward)하고 결과 계산
optimizer.zero_grad()
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
loss = criterion(outputs, labels)
# 역전파를 통해 기울기(gradient) 계산 및 학습 진행
loss.backward()
optimizer.step()
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
epoch_loss = running_loss / len(train_datasets)
epoch_acc = running_corrects / len(train_datasets) * 100.
# 학습 과정 중에 결과 출력
print('#{} Loss: {:.4f} Acc: {:.4f}% Time: {:.4f}s'.format(epoch, epoch_loss, epoch_acc, time.time() - start_time))
실행시켜보면 아래와 같이 몇번째 epochs인지가 출력되고 0부터 시작하므로 49번의 epochs이 출력되면 학습이 종료된다.
Acc는 정확도를 의미하며, AI Hub에서 제공해준 데이터가 좋은 탓인지 높은 정확도를 확인할수 있다.

이제 아래 코드를 이용하여 학습된 모델이 얼마나 높은 수준의 정확도를 갖는지 test폴더의 이미지들을 대상으로 확인할수 있다.
# 모델 평가
model.eval()
start_time = time.time()
with torch.no_grad():
running_loss = 0.
running_corrects = 0
correct = 0
false = 0
i = 0 # 100개만 테스트 하기 위해
for inputs, labels in test_dataloader:
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
loss = criterion(outputs, labels)
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
# 한 배치의 첫 번째 이미지에 대하여 결과 시각화
#print(f'[예측 결과: {class_names[preds[0]]}] (실제 정답: {class_names[labels.data[0]]})')
# 예측 결과 / 실제 정답
if class_names[preds[0]] == class_names[labels.data[0]]:
correct += 1
else:
false += 1
imshow(inputs.cpu().data[0], title=f'{i} : [predict: {class_names[preds[0]]}/result: {class_names[labels.data[0]]}]')
epoch_loss = running_loss / len(test_datasets)
epoch_acc = running_corrects / len(test_datasets) * 100.
print('[Test Phase] Loss: {:.4f} Acc: {:.4f}% Time: {:.4f}s'.format(epoch_loss, epoch_acc, time.time() - start_time))
print(f'correct : {correct} false : {false}')
실행 시키면 왼쪽에는 예측된 값을 오른쪽에는 실제 데이터의 값을 출력한다.
test이미지가 너무 많은 관계로, 이 과정은 test이미지를 따로 분비하여 해당 이미지를 대상으로 진행하기로 한다.

구글 드라이브에 model_test_image폴더를 만들고 시험해볼 이미지들을 넣은뒤, 이 이미지들을 대상으로 정확도와 함께 결과를 출력하는 파이썬 코드를 작성한다.
# 모델 정확도 검증
from PIL import Image
import torch.nn.functional as F
# model_test_images에 저장한 이미지들의 이름을 불러온다
folder_dir = r'./drive/MyDrive/model_test_images'
image_lists = os.listdir(folder_dir)
print(image_lists)
# 불러온 이미지 수만큼 모델로 테스트하여 정확도를 검정한다.
for i in range(len(image_lists)):
image = Image.open('./drive/MyDrive/model_test_images/' + image_lists[i])
image = transforms_test(image).unsqueeze(0).to(device)
model.to(device) # 모델을 동일한 디바이스로 이동시킴
with torch.no_grad():
outputs = model(image)
probabilities = F.softmax(outputs, dim=1) # 예측 결과를 확률로 변환
_, preds = torch.max(outputs, 1)
predicted_class = class_names[preds[0]]
predicted_prob = probabilities[0, preds[0]].item() # 예측된 클래스의 확률 값
imshow(image.cpu().data[0], title=f'predict_result: {predicted_class} (percentage: {predicted_prob:.2f})')
코드를 실행시키면 각 이미지에 대해 예측값과 확률을 출력한다.

이제 model에 resnet18로 학습된 모델이 담겼다.
우리는 이 model을 코랩에서만 사용할 수 없기 때문에 모델을 다운로드하여, 서버에서 작동시켜야한다.
모델을 저장하기 위한 코드는 아래와 같으며, 경로를 따로 설정하지 않는다면 코랩에 다운되기 때문에 코랩에 저장된 파일을 다시 데스크탑 등에 백업해야 한다. 6개의 모델로 학습되었다는 의미로 'resnet18_test3_model6.pth'로 저장하였다.
학습 모델은 주로 .pt또는 .pth라는 확장자로 저장된다.
import torch
# 모델 저장 => 구글 코랩에 다운되므로, 백업을 해야한다.
torch.save(model.state_dict(), 'resnet18_test3_model6.pth')
이제 만들어진 모델을 다시 불러와서 이전의 실행결과와 동일한 값을 같는지 확인해야한다.
같은 값을 갖는다면 모델을 제대로 불러왔다는 의미이다.
우리는 기존에 resnet18모델을 불러와서 전이학습을 거쳤고 그 과정에서 마지막 레이어의 출력 뉴런수를 6으로 변경하였기 때문에
모델을 정상적으로 이용하기 위해선 불러온 모델의 마지막 출럭 뉴런도 똑같이 6으로 변경해야한다.(객체에 대한 참조를 가져올때, 같은 타입의 객체로 받아야 하는것과 같다.)
resnet18의 경우 마지막 레이어가 512이므로 nn.Linear(512,6)을 통해 resnet18_test3_model6.pth파일과 동일한 상태로 만들어준다.
import torch
import torch.nn.functional as F
# 모델 매개변수 로드
model_parameters = torch.load('resnet18_test3_model6.pth')
# 새로운 모델 생성 => resnet18로 학습되었으므로, 기반 모델로 18로
model = torchvision.models.resnet18()
# 마지막 층의 출력 뉴런 수 변경 => 512는 resnet18의 마지막 레이어의 번호를 의미(마지막 레이어에서 뉴런 6개 이용?한것이므로)
model.fc = nn.Linear(512, 6) # 6은 학습한 객체의 수를 의미
# 모델 매개변수 로드
model.load_state_dict(model_parameters)
# 모델을 평가 모드로 설정
model.eval()
정상적으로 코드가 실행되면 아래와 같이 출력된다.

불러온 모델로 앞서 만든 모델 검증용 코드를 실행시켜, 같은 수준의 정확도를 갖는지 확인한다.

높은 정확도로 정상적으로 이미지가 분류되는것을 확인 할 수 있다.