추론 한 번에 차량 모델·색상까지 예측하기

차량 인식 시스템을 만들다 보면 이런 욕심이 자연스럽게 생깁니다.

• 이 이미지에서
  • 차량 대분류(승용 / RV / 트럭…),
  • 모델(소나타 / 그랜저…),
  • 색상(흰색 / 검정 / 회색…)
    을 한 번에 뽑아낼 수 없을까?

모델을 태스크별로 하나씩 따로 만들면

• 학습 스크립트도 여러 개,
• 배포할 모델도 여러 개,
• 추론할 때마다 이미지도 여러 번 넣어야 해서

확실히 번거롭습니다.

이걸 깔끔하게 정리해주는 방법이 바로 멀티태스크 러닝(Multi-Task Learning) 입니다.
실제 코드로 구현하면 “멀티헤드 분류(multi-head classification)” 구조에 가깝고요.

아래에서는 EfficientNet을 백본으로 쓰고,
large / medium / small 세 가지 분류를 한 번에 예측하는 구조와 학습 코드를 같이 보면서
실제로 어떻게 쓰는지 정리해보겠습니다.

---

멀티태스크 러닝, 개념부터 가볍게 짚고 가기

우리가 하고 싶은 건 한 줄로 요약하면 이겁니다.

이미지 한 장 넣고,
차량 대분류 / 중분류 / 소분류를 한 번에 예측한다.

이때 구조는 이렇게 나뉩니다.

• 앞쪽: 공통 백본(backbone)
  • EfficientNet 같은 CNN이 이미지에서 특징(feature)을 뽑는 부분
• 뒤쪽: 태스크별 헤드(head)
  • 대분류용 FC
  • 중분류용 FC
  • 소분류용 FC

즉, “특징 추출은 하나, 출력층은 여러 개” 인 구조입니다.

이 방식을 쓰면:

• 공통 특징을 공유해서 학습하니 과적합이 줄어들고
• GPU에서 한 번 forward로 여러 결과를 얻을 수 있어서 효율적이고
• 모델도 한 개만 관리하면 되니 운영이 편해집니다.

---

EfficientNet 백본 + 멀티헤드 구조

모델 쪽은 보통 이런 생각으로 설계합니다.

• EfficientNet 마지막 글로벌 풀링까지는 그대로 쓰고
• 원래 분류용 FC 레이어는 빼고
• 대신 태스크별로 FC 레이어를 따로 달아준다

개념만 보이도록 아주 단순화하면 대략 이런 느낌입니다.

import torch.nn as nn
from efficientnet_pytorch import EfficientNet  # 예시

class EfficientNetMultiTask(nn.Module):
    def __init__(self, num_large_classes, num_medium_classes, num_small_classes, pretrained=True):
        super().__init__()
        
        # 1) EfficientNet 백본 불러오기
        if pretrained:
            self.backbone = EfficientNet.from_pretrained('efficientnet-b0')
        else:
            self.backbone = EfficientNet.from_name('efficientnet-b0')
        
        in_features = self.backbone._fc.in_features
        # 원래 분류용 FC는 제거하고 feature extractor만 쓰도록 변경
        self.backbone._fc = nn.Identity()
        
        # 2) 태스크별 헤드
        self.fc_large = nn.Linear(in_features, num_large_classes)
        self.fc_medium = nn.Linear(in_features, num_medium_classes)
        self.fc_small = nn.Linear(in_features, num_small_classes)

    def forward(self, x):
        feat = self.backbone(x)  # 공통 feature
        
        out_large = self.fc_large(feat)
        out_medium = self.fc_medium(feat)
        out_small = self.fc_small(feat)
        
        # dict로 묶어 반환
        return {
            'large': out_large,
            'medium': out_medium,
            'small': out_small,
        }

여기서 중요한 포인트만 정리하면:

• EfficientNet은 “특징만 뽑는 역할”
• fc_large, fc_medium, fc_small 에서 각 태스크별로 분류
• forward() 는 dict 형태로 여러 출력이 나오게 설계

이렇게 만들어 두면 학습 루프에서 outputs['large'], outputs['medium'], outputs['small'] 이런 식으로 깔끔하게 다룰 수 있습니다.

---

멀티태스크 학습의 핵심: 태스크별 loss를 어떻게 합칠까?

모델 구조보다 더 중요한 부분이 바로 “loss를 어떻게 계산하느냐” 입니다.
멀티태스크에서는 보통 이렇게 합니다.

• 각 태스크마다 CrossEntropyLoss 계산
• 그걸 더하거나, 가중치를 곱해서 합산
• 그 합친 loss로 한 번 backward

여기서 사용한 train_epoch() 함수의 핵심 부분만 가져와 보면 흐름이 딱 보입니다.

def train_epoch(model, dataloader, criterion, optimizer, device, scaler=None):
    model.train()
    total_loss = 0.0
    large_correct = 0
    medium_correct = 0
    small_correct = 0
    total_samples = 0
    
    for batch in tqdm(dataloader, desc='Training'):
        images = batch['image'].to(device, non_blocking=True)
        large_labels = batch['large'].to(device, non_blocking=True)
        medium_labels = batch['medium'].to(device, non_blocking=True)
        small_labels = batch['small'].to(device, non_blocking=True)
        
        if scaler is not None:
            # AMP 사용 시
            with torch.cuda.amp.autocast():
                outputs = model(images)
                
                # 태스크별 loss
                loss_large = criterion(outputs['large'], large_labels)
                loss_medium = criterion(outputs['medium'], medium_labels)
                loss_small = criterion(outputs['small'], small_labels)
                
                # 동일 가중치로 합산
                loss = loss_large + loss_medium + loss_small
            
            optimizer.zero_grad()
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        else:
            # AMP 미사용 시
            outputs = model(images)
            
            loss_large = criterion(outputs['large'], large_labels)
            loss_medium = criterion(outputs['medium'], medium_labels)
            loss_small = criterion(outputs['small'], small_labels)
            loss = loss_large + loss_medium + loss_small
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        
        # 통계 계산
        total_loss += loss.item()
        total_samples += images.size(0)
        
        _, large_pred = torch.max(outputs['large'], 1)
        _, medium_pred = torch.max(outputs['medium'], 1)
        _, small_pred = torch.max(outputs['small'], 1)
        
        large_correct += (large_pred == large_labels).sum().item()
        medium_correct += (medium_pred == medium_labels).sum().item()
        small_correct += (small_pred == small_labels).sum().item()
    
    avg_loss = total_loss / len(dataloader)
    large_acc = large_correct / total_samples
    medium_acc = medium_correct / total_samples
    small_acc = small_correct / total_samples
    
    return avg_loss, large_acc, medium_acc, small_acc

여기서 볼 수 있는 멀티태스크 포인트는 딱 세 가지입니다.

1. 데이터셋이 image, large, medium, small 라벨을 한 번에 내보내고
2. 모델도 outputs['large'], ['medium'], ['small'] 이렇게 여러 출력을 내고
3. 각 태스크별 loss를 구한 뒤 loss = loss_large + loss_medium + loss_small 로 묶어서 역전파

원한다면 태스크 중요도에 따라

loss = 1.0 * loss_large + 0.5 * loss_medium + 0.5 * loss_small

처럼 가중치를 다르게 줄 수도 있습니다. 예를 들어 “대분류가 제일 중요하다” 같은 요구가 있을 때요.

검증 단계(validate)도 구조는 거의 같고, 대신 model.eval() + torch.no_grad() 로만 감싸서 쓰면 됩니다.

---

Mixed Precision, 멀티 GPU는 그냥 옵션으로 덧붙인 것

코드를 보면 AMP, 멀티 GPU 같은 옵션도 같이 들어가 있는데, 이건 말 그대로 “성능 튜닝용 껍데기”에 가깝습니다.

• torch.cuda.amp.autocast(), GradScaler()
  • H100 같은 GPU에서 계산을 더 빠르고, 메모리 적게 쓰게 해주는 옵션
• nn.DataParallel(model)
  • 여러 GPU 있을 때 배치를 나눠서 올려주는 방식

멀티태스크 러닝의 본질은 “백본 + 여러 헤드 + loss 합산” 이고,
AMP나 DataParallel은 그 위에 얹을 수 있는 가속/병렬화 옵션이라고 보면 됩니다.

---

실제 프로젝트에서 어떻게 확장하면 좋을까?

지금 코드는 large / medium / small 세 단계 분류를 한 번에 학습하는 예시인데,
같은 패턴으로 얼마든지 확장할 수 있습니다.

예를 들어:

• 차량 색상(head 하나 추가)
• 차량 진행 방향(front / rear / side)
• 차량에 특정 장비(루프박스, 택시 표시 등) 부착 여부

이런 것들을 전부 “추가 헤드”로 붙이면 됩니다.

구조는 그대로 유지하면서,
outputs 딕셔너리에 키만 하나씩 늘려가는 느낌입니다.

---

정리

• EfficientNet 같은 모델을 백본으로 쓰면
  “이미지 → 공통 특징” 까지는 한 번에 처리하고
• 그 뒤에 태스크별 FC 레이어를 올려서
  차량 대분류 / 중분류 / 소분류 / 색상 … 을 한 번에 예측할 수 있습니다.
• 학습에서는 태스크별 CrossEntropyLoss를 구하고,
  그걸 적당히 합산해서 한 번 backward 해주는 것이 멀티태스크 러닝의 핵심입니다.

지금 구조 그대로 가져다 쓰되,
태스크만 바꿔서 “컬러 + 모델 + 방향 + 번호판 여부” 같은 식으로 확장하면
현실적인 차량 인식 파이프라인을 꽤 깔끔하게 하나의 모델로 묶을 수 있습니다.