YOLO 클래스 수 조정과 설정법

숫자변경 사진


YOLO 기반의 객체 탐지 모델을 실제 프로젝트에 적용하면서 가장 많이 조정하게 되는 부분 중 하나가 바로 클래스 수 설정입니다. 처음에는 기본 제공되는 COCO 데이터셋 기반 구조를 그대로 사용하는 것이 편리하게 느껴졌지만, 실제 환경에서는 필요한 클래스만 남기거나 새로운 클래스를 추가하는 작업이 필수적이었습니다. 이 글에서는 YOLO에서 클래스 수를 변경하는 방법과 그에 따른 설정 파일 수정, 주의사항 등을 중심으로 커스텀 클래스를 적용하는 전체 과정을 정리하겠습니다.

YOLO에서 클래스 수가 중요한 이유

YOLO는 이미지 내에서 객체의 위치와 함께 클래스까지 동시에 예측하는 모델입니다. 따라서 모델의 출력층은 클래스 수에 따라 결정됩니다. 예를 들어, COCO 데이터셋을 기준으로 사전 학습된 YOLO 모델은 80개의 클래스를 구분할 수 있도록 구성되어 있습니다. 하지만 사용자의 목적에 따라 1개에서 수십 개까지 클래스 수를 조정하여 사용할 수 있으며, 이 경우 모델 구조, 설정 파일, 데이터 구성 모두에 변경이 필요합니다.

클래스 수가 변경되면 출력층의 차원이 달라지므로, 학습을 다시 시작해야 하며 사전 학습된 가중치를 그대로 사용할 수 없습니다. 다만 백본 가중치는 전이 학습 개념으로 사용할 수 있으므로, 기존 모델을 활용한 효율적인 학습도 가능합니다. 클래스 수가 많은 경우에는 다중 분류로 인한 정확도 저하가 발생할 수 있고, 클래스 간 유사성이 클 경우 오분류도 늘어날 수 있으므로, 클래스 설계 단계에서부터 충분한 고려가 필요합니다.

클래스 수 조정에 따른 설정 변경

YOLO에서 클래스 수를 변경하기 위해서는 몇 가지 필수 설정을 수정해야 합니다. 가장 먼저 확인해야 할 것은 데이터 구성 파일입니다. YOLOv5 기준으로 `data.yaml` 파일에 클래스 수와 클래스 이름이 명시되어야 합니다. 예를 들어 클래스가 3개라면, nc 항목을 3으로 지정하고, names 리스트에 해당 클래스 이름을 차례로 입력합니다.

예시 구성은 다음과 같습니다. train: ./images/train val: ./images/val nc: 3 names: [‘person’, ‘car’, ‘bicycle’] 이처럼 클래스 수(nc)와 이름(names)은 모델이 출력층에서 어떤 클래스를 예측해야 하는지를 정의합니다. 이 값을 정확히 지정하지 않으면, 모델은 올바른 출력을 생성하지 못하며 학습 또는 추론 시 오류가 발생할 수 있습니다.

다음으로는 모델 구성 파일의 수정이 필요합니다. YOLOv5에서는 `models/yolov5s.yaml`과 같은 모델 구조 정의 파일을 사용합니다. 해당 파일의 마지막 출력 부분에서 `nc` 값을 수정하여 클래스 수에 맞게 출력 채널을 조정해야 합니다. 예를 들어 기존에 `nc: 80`으로 설정되어 있다면, 이를 `nc: 3`으로 변경해야 합니다. 이 설정은 모델의 detection head가 몇 개의 클래스를 예측할지 결정하는 부분이므로, 꼭 잊지 않고 수정해야 합니다.

또한 기존 가중치를 사용할 경우에는 `--weights` 옵션으로 백본만 초기화하고, 출력층은 무작위로 초기화되도록 설정하는 것이 좋습니다. 이를 통해 불필요한 클래스 가중치를 제거하고 커스텀 클래스에 맞는 학습을 시작할 수 있습니다.

클래스 수 조정 시 주의사항

클래스 수를 줄이거나 늘릴 때에는 몇 가지 주의할 점이 있습니다. 첫째, 라벨 파일의 클래스 인덱스와 `data.yaml` 파일의 names 순서가 정확히 일치해야 합니다. 예를 들어 라벨 파일에서 ‘car’ 클래스가 1번이라면, `data.yaml`의 names 리스트에서도 1번 인덱스가 ‘car’여야 합니다. 이 순서가 어긋나면 모델은 잘못된 클래스를 학습하거나 예측하게 됩니다.

둘째, 클래스 수가 매우 적은 경우, 예를 들어 단일 클래스일 때에는 NMS 설정도 조정할 필요가 있습니다. 다중 클래스 탐지를 전제로 한 NMS 알고리즘은 클래스 구분 없이 박스를 제거하는 방식이기 때문에, 단일 클래스에서는 불필요한 제거가 발생할 수 있습니다. 이런 경우 `agnostic-nms` 설정을 확인하거나, post-processing에서 클래스 간 구분 여부를 조정해야 합니다.

셋째, 클래스 수가 많아질수록 데이터 균형이 중요해집니다. 특정 클래스에 데이터가 몰려 있다면, 모델은 해당 클래스에 편향되어 학습될 수 있습니다. 따라서 클래스별 이미지 수와 라벨 수를 미리 분석하여, 불균형이 있다면 증강이나 오버샘플링을 통해 보완하는 것이 좋습니다.

마지막으로 추론 시에도 `--classes` 옵션을 통해 특정 클래스만 추론하도록 설정할 수 있습니다. 이는 다중 클래스 모델을 학습했지만 실제로는 일부 클래스만 필요할 때 유용하게 사용할 수 있는 기능입니다. 반대로 전체 클래스를 대상으로 추론해야 하는 경우, 해당 옵션을 생략하여 전체 결과를 출력할 수 있습니다.

결론 – 클래스 수 조정은 YOLO 커스터마이징의 핵심

YOLO에서 클래스 수를 조정하는 과정은 단순히 숫자를 바꾸는 작업이 아니라, 모델 구조와 데이터 설정, 학습 전략 전반을 고려해야 하는 중요한 단계입니다. 정확한 클래스 정의와 설정 변경 없이 학습을 시작하게 되면, 탐지 결과가 왜곡되거나 모델이 정상적으로 수렴하지 않을 수 있습니다.

제가 진행했던 커스텀 탐지 프로젝트에서도 클래스 수 조정에 대한 실수가 여러 번 있었습니다. 초기에 라벨과 설정 파일 간의 불일치로 인해 학습된 모델이 엉뚱한 클래스를 예측하거나, 추론 결과가 전혀 나타나지 않는 문제가 발생했습니다. 이를 통해 클래스 수 조정이 단순한 설정 변경이 아닌, YOLO 모델 전체 동작에 영향을 주는 핵심적인 작업이라는 것을 실감했습니다.

YOLO를 활용한 객체 탐지 프로젝트를 안정적으로 수행하고자 한다면, 클래스 수 조정과 그에 따른 설정 변경을 체계적으로 수행하는 것이 중요합니다. 이 글이 클래스 수를 조정하고, 커스텀 모델을 구성하는 데 실질적인 도움이 되기를 바랍니다.

댓글

이 블로그의 인기 게시물

모델 평가 지표 정리 – 정확도, 정밀도, F1 Score까지

선형회귀와 로지스틱 회귀 – 가장 기초적인 예측 모델들

지도학습 vs 비지도학습 – 머신러닝의 대표 학습 방식 비교