250926 [ Day 59 ] - Project (3)

시작하며

모델을 학습시키기 위해서는 데이터 폴더 구조가 깔끔하게 정리되어 있어야 하는데, 구한 데이터들은 쓸데없는 폴더들이 너무 많아서 폴더 구조 정리가 시급했다. 하지만 수작업을 하기에는 폴더가 적어도 몇천 개는 넘었기에 폴더 구조를 통일할 수 있는 코드를 구상하면서 시작했다.

Folder Split

import os, shutil, glob

# 원본 폴더 경로 지정
source_path = "./path/"

# 클래스 이름 설정
def folder_split(source_path, class_name="class_name"):

    for folder in ["train", "valid", "test"]:
        # 원본 폴더 경로 지정
        source_split_path = os.path.join(source_path, folder)

        # 타겟 경로 설정 및 폴더 생성
        if os.path.exists(source_split_path):
            target_dir = os.path.join(source_path, class_name, folder)
            os.makedirs(target_dir, exist_ok=True)

            # 이미지 탐색
            images = glob.glob(os.path.join(source_split_path, "**", "*.jpg"), recursive=True)
            images += glob.glob(os.path.join(source_split_path, "**", "*.png"), recursive=True)

            # 이미지 복사
            for idx, img_path in enumerate(images):
                # 파일 확장자 추출
                file_ext = os.path.splitext(img_path)[1]
                # 새로운 파일명 생성
                new_filename = f"img_{idx:06d}{file_ext}"
                # 새로운 파일 경로 정리
                target_path = os.path.join(target_dir, new_filename)
                # 파일 복사
                shutil.copy2(img_path, target_path)
        else:
            print(f"{source_split_path} 폴더를 찾을 수 없습니다.")

    print("완료!")

# 파일 실행
folder_split(source_path)

전부터 느꼈지만 OS를 잘 활용하면 유용한 부분이 참 많은 것 같다. 다음에 OS 활용에 대해 깊이 공부해 보고 싶다는 생각이 들었다.

Model Test

폴더 구조 정리가 끝나고 간단하게 몇 개의 데이터를 사용해서 AI 모델을 학습시켜 보았다. PyTorch 모델과 YOLO를 사용할 것 같은데, 먼저 PyTorch의 ResNet18 모델을 사용해 보았다.

import torch
import torchvision.models as models
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, datasets
from torch.utils.data import DataLoader

# 이미 학습되어 있는 resnet18 모델 불러오기
model = models.resnet18(pretrained=True)
model.fc = nn.Linear(512, 1)

# 데이터 불러오기
transform = transforms.Compose([
    # 데이터 전처리 (224x224) 사이즈로 통일
    transforms.Resize((224,224)),
    # 텐서로 변환
    transforms.ToTensor(),
    # ImageNet 통계로 정규화
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 데이터 불러오기
train_dataset = datasets.ImageFolder("dataset/data/train", transform=transform)
test_dataset = datasets.ImageFolder("dataset/data/test", transform=transform)
valid_dataset = datasets.ImageFolder("dataset/data/valid", transform=transform)

# 배치 사이즈 설정
batch_size = 32

# 데이터 로더 생성
train_loader = DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle=True, # 에포크마다 섞어서 훈련
    num_workers=0 # 병렬처리할 서브 프로세스 / 윈도우 multiprocessing error시 0으로 처리
)
test_loader = DataLoader(
    test_dataset,
    batch_size=batch_size,
    shuffle=False,
    num_workers=0
)
valid_loader = DataLoader(
    valid_dataset,
    batch_size=batch_size,
    shuffle=False,
    num_workers=0
)

# 사용가능한 디바이스 확인
if torch.cuda.is_available():
    device = torch.device('cuda')
    print("CUDA GPU 사용")
    print(f"GPU 이름: {torch.cuda.get_device_name(0)}")
elif torch.backends.mps.is_available():
    device = torch.device('mps')
    print("Apple MPS (Metal Performance Shaders) 사용")
else:
    device = torch.device('cpu')
    print("CPU 사용")

print(f"선택된 디바이스: {device}")

# 모델을 디바이스로 이동
model = model.to(device)

# 손실 함수 정의
# 테스트는 클래스1개 -> 이진분류 사용 추후 여러 클래스 활용시 Cross Entropy로 변경
criterion = nn.BCEWithLogitsLoss()

# 옵티마이저 설정
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 학습 루프 함수 정의
def train_loop(data_loader, model, criterion, optimizer):
    model.train() # 학습 모드
    size = len(data_loader.dataset)
    running_loss = 0.

    for batch_idx, (data, target) in enumerate(data_loader):
        device = next(model.parameters()).device
        data, target = data.to(device), target.float().to(device)

        optimizer.zero_grad()
        outputs = model(data)
        outputs = outputs.squeeze(-1) # 이진분류 처리
        loss = criterion(outputs, target)
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * data.size(0)

        if (batch_idx+1) % 10 == 0:
            current = batch_idx * len(data)
            print(f"[batch : {batch_idx+1: 4d}], Loss : {loss.item():>7f} ({current:>5d} / {size:>5d})")
    
    epoch_loss = running_loss / size
    return epoch_loss

# 검증 루프 함수 정의
def valid_loop(data_loader, model, criterion):
    model.eval() # 평가 모드
    size = len(data_loader.dataset)
    valid_loss = 0.
    correct = 0

    with torch.no_grad():
        for data, target in data_loader:
            device = next(model.parameters()).device
            data, target = data.to(device), target.float().to(device)

            outputs = model(data)
            outputs = outputs.squeeze(-1) # 이진분류 처리
            loss = criterion(outputs, target)
            valid_loss += loss.item() * data.size(0)

            # 정확도 계산
            predictions = torch.sigmoid(outputs) > 0.5
            correct += (predictions == target.bool()).sum().item()
        
    # 평균 계산
    avg_loss = valid_loss / size
    accuracy = correct / size

    print(f"Validation Results: Accuracy: {accuracy:.3f} ({100*accuracy:.1f}%), Avg loss: {avg_loss:.4f}")
    return avg_loss, accuracy

# 학습 실행
epochs = 5 # 반복 학습

for epoch in range(epochs):
    print(f"\n[Epoch] {epoch+1} / {epochs}")
    train_loss = train_loop(train_loader, model, criterion, optimizer)
    valid_loss, valid_acc = valid_loop(valid_loader, model, criterion)

print("\n학습 및 검증 완료!")

torch.save(model.state_dict(), "test_model.pth")

print("\n모델 저장 완료!")

Trouble Shooting

Windows Multiprocessing Error

num_workers 는 병렬 처리를 위한 서브 프로세스의 수로, 보통 CPU 코어 수에 맞추어 설정하지만, 오류 발생 시 0으로 처리

ValueError: Target size (torch.Size([1])) must be the same as input size (torch.Size([]))

배치의 차원 불일치가 원인으로 outputs = outputs.squeeze() 에서 마지막 차원만 제거하기 위해 outputs = outputs.squeeze(-1) 로 처리

Test 학습 및 검증 결과

빠른 테스트를 위해 Epoch 5 설정

• Epoch 1 = Acc : 90.9% / Avg Loss : 0.3099
• Epoch 2 = Acc : 96.0% / Avg Loss : 0.1112
• Epoch 3 = Acc : 93.8% / Avg Loss : 0.1488
• Epoch 4 = Acc : 93.6% / Avg Loss : 0.1627
• Epoch 5 = Acc : 97.1% / Avg Loss : 0.0788
  
  

마치며

세 번째 에포크부터 성능이 낮아져서 과적합이 생긴 건지 의심이 됐는데, 최종적으로는 97%까지 정확도가 올라서 일단 테스트 학습 및 검증은 성공적으로 마쳤다. 역시 기존에 학습이 되어 있는 모델을 사용하니까 정확도가 상당히 높게 나오는 것 같다. 내일은 오늘 저장한 모델을 어제 만들어 둔 테스트용 API에 연동시켜서, 웹에서 이미지를 업로드하고 AI 모델을 거쳐 결과를 도출하는 기능을 목표로 작업을 해 볼 예정이다.