Transfer Learning으로 소량 데이터 문제 해결하기: ImageNet 사전학습 CNN으로 산업 설비 고장 진단 구현하기

import librosa
import librosa.display
import matplotlib.pyplot as plt
import numpy as np

# 진동 신호를 Spectrogram으로 변환
def vibration_to_spectrogram(signal, sr=10000, save_path='spec.png'):
    # STFT 수행
    D = librosa.stft(signal)
    S_db = librosa.amplitude_to_db(np.abs(D), ref=np.max)

    # 시각화 및 저장
    plt.figure(figsize=(10, 4))
    librosa.display.specshow(S_db, sr=sr, x_axis='time', y_axis='hz')
    plt.colorbar(format='%+2.0f dB')
    plt.tight_layout()
    plt.savefig(save_path, bbox_inches='tight', pad_inches=0)
    plt.close()

import torch
import torch.nn as nn
from torchvision import models

# ResNet50 사전학습 모델 로드
model = models.resnet50(weights='IMAGENET1K_V2')

# 마지막 FC layer만 교체 (고장 유형 3개 분류 예시)
num_classes = 3  # 정상, 베어링 고장, 언밸런스
model.fc = nn.Linear(model.fc.in_features, num_classes)

# GPU 사용
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

# Partial Fine-tuning 예시: ResNet 앞 3개 layer만 freeze
for name, param in model.named_parameters():
    if 'layer1' in name or 'layer2' in name or 'layer3' in name or 'conv1' in name:
        param.requires_grad = False  # freeze
    else:
        param.requires_grad = True   # 학습 허용

# Optimizer 설정 (학습 가능한 파라미터만)
optimizer = torch.optim.Adam(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=0.0001
)

from torchvision import transforms

# 학습 데이터용 증강
train_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomRotation(15),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 검증/테스트 데이터용 (증강 없음)
val_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder

# 데이터 로드
train_dataset = ImageFolder('data/train', transform=train_transforms)
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)

val_dataset = ImageFolder('data/val', transform=val_transforms)
val_loader = DataLoader(val_dataset, batch_size=16, shuffle=False)

# 학습 루프
criterion = nn.CrossEntropyLoss()
num_epochs = 30

for epoch in range(num_epochs):
    model.train()
    train_loss = 0.0

    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()

    # 검증
    model.eval()
    val_correct = 0
    val_total = 0

    with torch.no_grad():
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            val_total += labels.size(0)
            val_correct += (predicted == labels).sum().item()

    val_acc = 100 * val_correct / val_total
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {train_loss/len(train_loader):.4f}, Val Acc: {val_acc:.2f}%')

# Weighted Loss 예시 (고장 클래스 가중치 높임)
class_weights = torch.tensor([1.0, 3.0, 3.0]).to(device)  # 정상:고장1:고장2
criterion = nn.CrossEntropyLoss(weight=class_weights)