더 작고 빠른 모델에서 고성능 달성하기

모델 경량화가 중요한 시대다. 특히 DAMO-YOLO처럼 정확도는 뛰어나지만 계산 비용이 큰 디텍션 모델을 모바일 환경에 배포하려면, 성능과 속도 사이의 균형이 필수적이다. 이 문제를 해결하는 핵심 기술 중 하나가 바로 **지식 증류**(Knowledge Distillation)다. 이 방식은 이미 훈련된 고성능 ‘교사 모델’이 가진 패턴 인식 능력을, 파라미터 수가 적은 ‘학생 모델’이 습득하도록 유도하는 전이 학습의 일종이다. 본 글에서는 DAMO-YOLO의 지식을 MobileNetV3 기반의 경량 네트워크로 효과적으로 전달하는 방법을 다룬다.

1. 개발 환경 구성

지식 증류를 구현하기 위해 필요한 핵심 도구들을 설치하자. 여기서는 TinyNAS 외에도 일반적인 딥러닝 및 컴퓨터비전 라이브러리를 활용한다.

pip install torch torchvision torchaudio
pip install tinynas opencv-python-headless
pip install onnx onnxruntime

GPU를 사용한다면 CUDA 11.8 이상과 호환되는 PyTorch 버전을 선택해야 한다. 또한, 모델의 추론 최적화와 모바일 배포를 염두에 두고 ONNX 변환 기능도 함께 준비한다. TinyNAS는 아키텍처 탐색뿐 아니라 사전 정의된 증류 파이프라인을 제공해, 초보자도 비교적 쉽게 적용할 수 있다.

2. 지식 증류의 핵심 개념

전통적인 분류 학습은 정답(label)에 대한 one-hot 인코딩을 기반으로 한다. 반면 지식 증류는 교사 모델의 출력 로짓(logit)에서 얻은 **소프트 확률 분포**를 추가 학습 신호로 사용한다.

이 방식의 장점은 다음과 같다:

• 교사 모델이 각 클래스 간 유사성을 내재적으로 포착한 정보를 학생 모델이 습득 가능
• 학습 데이터에 포함되지 않은 암묵적 지식(예: "개"와 "늑대"는 서로 더 가깝다) 전이 가능
• 매우 작은 학생 모델이라도 상당한 성능 유지 가능

증류 과정에서 교사 모델은 eval() 상태로 고정되며, 그래디언트 계산이 비활성화된다. 학생 모델만 파라미터 업데이트 대상이 된다.

3. 교사-학생 아키텍처 설정

아래 코드는 교사 모델(DAMO-YOLO)과 학생 모델(MobileNetV3-Large)을 로드하고, 출력 차원을 맞추기 위한 어댑터를 추가하는 예시다.

import torch
import torch.nn as nn
from damo_detector import create_damo_yolo
from mobilenetv3 import build_mobilenetv3

# 교사 모델 로드 (학습된 가중치 사용)
teacher = create_damo_yolo(num_classes=80, pretrained=True)
teacher.eval()

# 학생 모델 생성
student = build_mobilenetv3(num_classes=80)

# 출력 크기 조정을 위한 프로젝션 레이어 (필요 시)
if student.classifier.in_features != 80:
    student.classifier = nn.Linear(student.classifier.in_features, 80)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
teacher.to(device)
student.to(device)

출력 노드 수가 다르면 성능 저하가 발생할 수 있으므로, 분류 헤드의 차원을 반드시 동일하게 맞춰야 한다.

4. 복합 손실 함수 설계

지식 증류의 성패는 손실 함수 설계에 달려 있다. 아래는 하드 타겟(정답 레이블)과 소프트 타겟(교사 출력)을 결합한 손실 함수 구현 예이다.

import torch.nn.functional as F

def combined_distill_loss(y_student, y_teacher, y_true, alpha=0.8, T=5.0):
    # 실제 레이블에 대한 학습 (Hard Target Loss)
    loss_hard = F.cross_entropy(y_student, y_true)
    
    # 교사 모델의 출력을 따르는 학습 (Soft Target Loss)
    p_t = F.softmax(y_teacher / T, dim=1)
    p_s = F.log_softmax(y_student / T, dim=1)
    loss_soft = F.kl_div(p_s, p_t, reduction='batchmean') * (T * T)
    
    # 가중합
    total_loss = alpha * loss_soft + (1 - alpha) * loss_hard
    return total_loss

여기서 중요한 하이퍼파라미터는 다음과 같다:

• T (Temperature): 확률 분포의 부드러움 조절. 일반적으로 2~10 사이 사용.
• alpha: 소프트 손실의 비중. 초반에는 낮게 시작해 점진적으로 증가시키는 스케줄링 전략이 효과적.

5. 효율적인 훈련 전략

단순히 한 번에 전체 데이터로 훈련하는 것보다, 다음의 전략들이 성능 향상에 기여한다:

5.1 단계별 증류 (Curriculum Distillation)

초반 에폭에서는 단순 이미지나 기본 객체부터 학습시켜 안정적인 기반을 마련한 후, 점차 복잡한 샘플을 포함시킨다.

5.2 온도 스케줄링

훈련 초기에는 높은 T 값으로 일반화된 지식을 전이하고, 후반으로 갈수록 T를 낮춰 세밀한 특징을 전달한다.

5.3 학습률 스케줄링

학생 모델은 랜덤 초기화이므로, 교사 모델보다 더 높은 초기 학습률(예: 1e-3)을 설정하고 Cosine Annealing 등을 적용한다.

optimizer = torch.optim.AdamW(student.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

for epoch in range(100):
    for x_batch, y_batch in train_loader:
        x_batch, y_batch = x_batch.to(device), y_batch.to(device)
        
        with torch.no_grad():
            y_teacher = teacher(x_batch)
        
        y_student = student(x_batch)
        loss = combined_distill_loss(y_student, y_teacher, y_batch)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    scheduler.step()

6. 성능 평가 및 모바일 배포

증류 후 학생 모델의 성능을 평가하면, 일반적으로 교사 모델 대비 mAP 기준 90% 이상의 정확도를 유지하면서도, 파라미터 수는 1/5 수준, 추론 속도는 3~4배 빨라지는 결과를 얻을 수 있다.

배포 최적화 단계

1. ONNX 변환: 학습된 모델을 ONNX 형식으로 내보내 다양한 플랫폼에서 실행 가능하게 함.
2. 양자화: INT8 또는 FP16 양자화를 통해 모델 크기 감소 및 추론 가속화.
3. 모바일 프레임워크 통합: TensorFlow Lite, NCNN, MNN 등에 최적화하여 실제 기기에서 테스트.

# ONNX로 내보내기 예시
dummy_input = torch.randn(1, 3, 640, 640).to(device)
torch.onnx.export(
    student,
    dummy_input,
    "distilled_det.onnx",
    input_names=["input"],
    output_names=["output"],
    opset_version=13,
    dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}}
)

7. 마무리

DAMO-YOLO와 같은 고성능 디텍터를 모바일 환경에 배포하기 위해 지식 증류는 매우 실용적인 해법이다. MobileNetV3와 같은 경량 백본에 DAMO-YOLO의 인식 능력을 전이함으로써, 리소스 제약 환경에서도 실시간 객체 탐지가 가능해진다. 특히 TinyNAS와 같은 도구를 활용하면, 아키텍처 설계부터 증류까지의 전 과정을 자동화할 수 있어 개발 효율성이 크게 향상된다. 핵심은 손실 함수의 하이퍼파라미터 튜닝과 훈련 스케줄링에 있으며, 실험을 통해 최적의 조합을 찾는 것이 중요하다. 모바일 AI 서비스를 개발 중이고 모델 크기 문제로 고민이라면, 지식 증류를 실험해보는 것을 강력히 추천한다.