PyTorch로 만드는 AI 재활용 분류기
전이학습으로 캔, 플라스틱, 종이 구별하기
AI와 딥러닝 기술이 우리 삶에 점점 더 가까워지고 있습니다. 오늘은 딥러닝 기술을 활용해 우리 주변의 문제를 해결하는 흥미로운 프로젝트를 함께 진행해보려 합니다. 바로 AI를 이용해 재활용품을 자동으로 분류하는 'AI 재활용 분류기'를 만들어보는 것입니다. 이 프로젝트는 『실전! 프로젝트로 배우는 딥러닝 컴퓨터비전 (개정판)』의 5장 '이미지 분류를 활용한 재활용품 분류' 내용을 기반으로 합니다. 복잡한 이론에 앞서, 실제 작동하는 모델을 직접 만들어보면서 딥러닝의 매력에 빠져보시길 바랍니다.
특히 이번 튜토리얼에서는 '전이학습(Transfer Learning)'이라는 강력한 기법을 사용할 것입니다. 전이학습은 수백만 장의 이미지로 미리 훈련된 전문가급 모델을 활용하는 방법으로, 우리가 가진 상대적으로 적은 데이터만으로도 매우 높은 성능의 분류 모델을 만들 수 있게 해줍니다.
자, 그럼 PyTorch와 함께 AI 재활용 분류기 만들기를 시작해볼까요?
1. 데이터 준비: PyTorch 커스텀 데이터셋 클래스 구성하기
딥러닝 모델을 훈련시키려면, 먼저 데이터를 모델이 이해할 수 있는 형태로 공급해주어야 합니다. 여기서는 'can', 'glass', 'paper', 'plastic'과 같이 클래스별로 폴더가 정리된 이미지 데이터셋을 PyTorch에서 효율적으로 사용하기 위해 Dataset 클래스를 직접 만드는 방법을 알아봅니다. 이 클래스는 이미지 파일을 불러와 딥러닝 모델에 맞는 크기와 형식으로 변환하는 전처리 과정까지 담당합니다.
참조: 『실전! 프로젝트로 배우는 딥러닝 컴퓨터비전 (개정판)』 5장 2절, Dataset_Class.py
클래스: PyTorch_Classification_Dataset_Class
코드 분석
PyTorch의 Dataset 클래스를 상속받아 우리만의 데이터셋 클래스를 만들 때는 __init__, __len__, __getitem__ 세 가지 메서드가 핵심입니다.
Python
import os
from torch.utils.data import Dataset
import torchvision.transforms as transforms
from PIL import Image
class PyTorch_Classification_Dataset_Class(Dataset):
def __init__(self, dataset_dir, transform=None):
# 1. 데이터셋 경로와 전처리 방법을 초기화합니다.
self.image_abs_path = dataset_dir
self.transform = transform
# 2. 클래스(폴더명) 리스트를 생성하고 정렬합니다.
self.label_list = os.listdir(self.image_abs_path)
self.label_list.sort()
self.x_list = [] # 이미지 경로 리스트
self.y_list = [] # 라벨 인덱스 리스트
# 3. 모든 이미지 경로와 해당 라벨 인덱스를 리스트에 저장합니다.
for label_index, label_str in enumerate(self.label_list):
img_path = os.path.join(self.image_abs_path, label_str)
img_list = os.listdir(img_path)
for img in img_list:
self.x_list.append(os.path.join(img_path, img))
self.y_list.append(label_index)
def __len__(self):
# 4. 전체 이미지 개수를 반환합니다.
return len(self.x_list)
def __getitem__(self, idx):
# 5. idx에 해당하는 이미지와 라벨을 가져옵니다.
image = Image.open(self.x_list[idx])
# 6. 이미지가 흑백일 경우 RGB로 변환합니다.
if image.mode != "RGB":
image = image.convert('RGB')
# 7. 정의된 전처리를 이미지에 적용합니다.
if self.transform is not None:
image = self.transform(image)
return image, self.y_list[idx]
__init__(self, ...): 클래스가 생성될 때 실행됩니다. 데이터셋의 루트 폴더(dataset_dir)를 받아와 (2) 각 하위 폴더(can, glass 등)의 이름을 라벨로 인식하고, (3) 폴더 내 모든 이미지 파일의 전체 경로와 숫자 라벨(예: 'can' -> 0, 'glass' -> 1)을 각각 x_list와 y_list에 저장합니다.
__len__(self): 데이터 로더가 전체 데이터셋의 크기를 알 수 있도록 x_list의 길이, 즉 총 이미지의 개수를 반환합니다.
__getitem__(self, idx): 데이터 로더가 특정 인덱스(idx)의 데이터를 요청할 때 호출됩니다. (5) x_list에서 이미지 경로를 찾아 PIL 라이브러리로 이미지를 열고, (7) __init__에서 정의한 transform(리사이즈, 크롭, 텐서 변환, 정규화 등)을 적용한 뒤, 이미지 텐서와 숫자 라벨을 함께 반환합니다.
2. 모델 준비: 전이학습으로 MobileNet V2 활용하기
강력한 이미지 분류 모델을 처음부터 만드는 것은 엄청난 데이터와 시간이 필요합니다. '전이학습'은 이 문제를 해결하는 아주 효율적인 방법입니다. 이미 수백만 장의 이미지로 훈련을 마친 전문가급 모델(사전 훈련된 모델)을 가져와, 우리 프로젝트의 목적에 맞게 마지막 부분만 살짝 수정하여 사용하는 기법입니다. 이 책에서는 가볍고 성능이 뛰어난 MobileNet V2 모델을 활용합니다.
코드 위치
참조: 『실전! 프로젝트로 배우는 딥러닝 컴퓨터비전 (개정판)』 5장 4절, Model_Class_Transfer_Learning_MobileNet.py
클래스: MobileNet
코드 분석
Python
import torch.nn as nn
import torchvision.models as models
class MobileNet(nn.Module):
def __init__(self, num_classes, pretrained=True):
super().__init__()
# 1. ImageNet으로 사전 훈련된 MobileNet V2 모델을 불러옵니다.
self.network = models.mobilenet_v2(pretrained=pretrained)
# 2. 모델의 마지막 분류 레이어의 입력 피처 수를 가져옵니다.
num_ftrs = self.network.classifier[1].in_features
# 3. 마지막 분류 레이어를 우리 프로젝트에 맞게 새로운 레이어로 교체합니다.
self.network.classifier[1] = nn.Linear(num_ftrs, num_classes)
self.classifier = nn.Sequential(nn.Softmax(dim=-1))
def forward(self, x):
x = self.network(x)
x = self.classifier(x)
return x
models.mobilenet_v2(pretrained=True): 이 한 줄이 전이학습의 핵심입니다. pretrained=True 옵션은 PyTorch가 ImageNet 데이터셋으로 미리 훈련시켜놓은 MobileNet V2의 구조와 가중치를 그대로 다운로드하라는 의미입니다. 이 모델은 이미 이미지의 기본적인 특징(선, 면, 질감 등)을 학습한 상태입니다.
num_ftrs = ...: 기존 MobileNet V2의 마지막 분류 레이어(classifier[1])가 입력으로 받는 특징(feature)의 개수를 알아냅니다.
self.network.classifier[1] = ...: 가장 중요한 부분입니다. 기존의 1000개 클래스 분류 레이어를 제거하고, 동일한 입력 피처 수(num_ftrs)를 받아서 우리에게 필요한 num_classes(재활용품 4종) 만큼의 출력값을 내보내는 새로운 nn.Linear 레이어로 교체합니다. 이로써 전문가 모델을 우리만의 '재활용 분류기'로 맞춤 제작하게 됩니다.
3. 딥러닝 모델 훈련시키기
데이터와 모델이 준비되었으니, 이제 훈련을 통해 모델을 '학습'시킬 차례입니다. 훈련은 모델에게 재활용품 이미지를 반복해서 보여주고, 정답을 맞히도록 가중치(weight)를 조금씩 조정해나가는 과정입니다. 이 과정은 훈련 데이터로 학습하고, 검증 데이터로 성능을 평가하며 최적의 모델을 찾아 저장하는 방식으로 진행됩니다.
참조: 『실전! 프로젝트로 배우는 딥러닝 컴퓨터비전 (개정판)』 5장 5절, Training_Class.py
메서드: training_network
코드 분석 (핵심 로직)
Python
# (Training_Class 내의 training_network 메서드 핵심부)
for epoch in range(1, epochs + 1):
# 1. 모델을 훈련 모드로 설정
self.model.train()
for data, target in tqdm(self.train_loader):
data, target = data.to(self.DEVICE), target.to(self.DEVICE)
# 2. 이전 배치의 기울기(gradient) 초기화
optimizer.zero_grad()
# 3. 순전파: 모델에 데이터를 입력하여 예측값(output) 계산
output = self.model(data)
# 4. 손실 계산: 예측값과 실제 정답(target)을 비교하여 오차(loss) 계산
loss = F.cross_entropy(output, target)
# 5. 역전파: 손실을 기반으로 각 가중치의 기울기 계산
loss.backward()
# 6. 가중치 업데이트: 계산된 기울기를 이용해 모델의 가중치 수정
optimizer.step()
# --- 한 Epoch 훈련 후 검증 시작 ---
self.model.eval() # 7. 모델을 평가 모드로 설정
with torch.no_grad(): # 8. 기울기 계산 비활성화
for data, target in tqdm(self.test_loader):
# ... 검증 데이터로 손실 및 정확도 계산 ...
# ... 가장 성능이 좋은 모델 저장 ...
model.train(): 모델을 '훈련 모드'로 전환합니다. 이는 드롭아웃(Dropout)이나 배치 정규화(Batch Normalization) 같은 훈련 시에만 필요한 기능들을 활성화합니다.
optimizer.zero_grad(): PyTorch는 기울기 값을 계속 누적하기 때문에, 새로운 배치를 처리하기 전에 항상 0으로 초기화해야 합니다.
output = self.model(data) (순전파): 전처리된 이미지 데이터(data)를 모델에 통과시켜 각 클래스에 대한 확률 예측값(output)을 얻습니다.
loss = F.cross_entropy(...) (손실 계산): 예측값과 실제 정답(target)을 비교하여 모델이 얼마나 틀렸는지 오차(손실)를 계산합니다.
loss.backward() (역전파): 계산된 손실을 기반으로, 모델의 각 가중치가 손실에 얼마나 영향을 미쳤는지(기울기) 역방향으로 계산합니다.
optimizer.step() (가중치 업데이트): 계산된 기울기를 바탕으로 옵티마이저(예: Adam)가 모델의 가중치를 손실이 줄어드는 방향으로 조금씩 업데이트합니다. 이 과정이 바로 '학습'입니다.
model.eval(): 검증 시에는 모델을 '평가 모드'로 전환하여 불필요한 기능(드롭아웃 등)을 비활성화합니다.
with torch.no_grad(): 기울기를 계산할 필요가 없는 검증 과정에서는 이 구문을 사용하여 불필요한 연산을 막고 메모리를 절약합니다.
마치며
지금까지 PyTorch를 이용해 AI 재활용 분류기를 만드는 전체 과정을 단계별로 살펴보았습니다. 커스텀 데이터셋을 구성하고, 전이학습을 통해 강력한 MobileNet V2 모델을 우리 프로젝트에 맞게 수정한 뒤, 훈련과 검증을 통해 최종 모델을 완성하는 흐름을 이해하셨을 것입니다.
이 블로그 글은 프로젝트의 큰 그림을 이해하는 데 초점을 맞췄습니다. 각 단계의 상세한 코드 구현, 이미지와 동영상 파일을 이용한 실제 추론 방법, 그리고 완성된 모델을 NVIDIA 젯슨 나노와 같은 임베디드 장치에 배포하는 실전 노하우까지, 프로젝트의 모든 것은 『실전! 프로젝트로 배우는 딥러닝 컴퓨터비전 (개정판)』에 담겨 있습니다.
지금까지 살펴본 내용을 바탕으로 여러분만의 AI 프로젝트에 도전해보시는 것은 어떨까요? 이 글이 그 첫걸음에 작은 도움이 되었기를 바랍니다.
https://wikibook.co.kr/dlvision-project-rev/
딥러닝 영상처리의 핵심 이론과 컴퓨터비전·자율주행 프로젝트를 배웁니다! 이 책은 딥러닝 컴퓨터비전과 자율 주행의 전문가가 되고자 하는 대학생과 개발 입문자를 위한 학습 자료로, 딥러닝 영상처리의 핵심 이론과 관련 문제 해결에 적용 가능한 입문용 미니 프로젝트들을 다룬다
