신경망(Neural Networks): PyTorch를 활용한 전이학습 기반 이미지 불량/정상 분류 모델 예제

"지도 학습(supervised learning) > 분류(classification) > 이진 분류(binary classification) > 신경망(Neural Network, 딥러닝)"에 해당하는 예제입니다.
※ "스팸 필터"와 역할(작동 방식)이 동일(특성 → 정답 맞추기)하되, 분류 대상이 "텍스트"가 아니라 "이미지(비전)"라는 차이만 있습니다.

파이썬 패키지 설치(필요시)

# !pip install torch torchvision matplotlib tqdm

1) 필요한 라이브러리 임포트

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader
from tqdm.notebook import tqdm  # 진행률 표시 (optional)

2) 데이터 전처리 및 데이터로더 구성

# 이미지 전처리 및 데이터셋 생성
transform = transforms.Compose([
    transforms.Resize((224, 224)),   # ResNet 입력 사이즈
    transforms.ToTensor(),
    # transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # 필요시 활성화
])

train_dataset = datasets.ImageFolder('data/train', transform=transform)
val_dataset   = datasets.ImageFolder('data/val',   transform=transform)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader   = DataLoader(val_dataset, batch_size=32, shuffle=False)

Tip:
data/train/정상/이미지A.jpg, data/train/불량/이미지B.jpg 폴더 구조로 준비되어 있어야 합니다.

https://ndb796.tistory.com/658#google_vignette 이미지 없으니까 여기이 이거어

3) 사전학습(Pretrained) 모델 불러오기 및 수정

# 사전학습된 ResNet18 불러오기
model = models.resnet18(pretrained=True)
num_classes = 2  # 정상/불량
model.fc = nn.Linear(model.fc.in_features, num_classes)

4) 손실함수/옵티마이저 정의

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-4)

5) 학습 & 검증 루프

# CUDA 사용 가능 여부
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

num_epochs = 10  # 에폭수 조정 가능

for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, labels in tqdm(train_loader, desc=f"[Epoch {epoch+1}/Train]"):
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * images.size(0)
    avg_loss = running_loss / len(train_loader.dataset)
    print(f"Epoch {epoch+1}, Average Train Loss: {avg_loss:.4f}")

    # Validation
    model.eval()
    correct, total = 0, 0
    with torch.no_grad():
        for images, labels in tqdm(val_loader, desc=f"[Epoch {epoch+1}/Val]"):
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            correct += (predicted == labels).sum().item()
            total += labels.size(0)
    acc = 100 * correct / total
    print(f"Validation Accuracy: {acc:.2f}%\n")

6) 모델 저장

torch.save(model.state_dict(), 'inspection_model.pth')
print("Model saved as inspection_model.pth")

완료 후 체크포인트

• GPU 사용 여부는 torch.cuda.is_available()로 자동 적용 (GPU 있으면 훨씬 빠름)

• 폴더 구조 및 이미지 준비상태 확인

• 실제 데이터 수가 적으면 overfitting 가능 (따라서 실습 후엔 augmentation 등 추가 권장)