커스텀 데이터셋을 활용한 시맨틱 세그멘테이션을 위한 SegFormer 파인튜닝

Hugging Face는 커스텀 이미지 데이터셋을 처리하기 위해 최첨단 시맨틱 세그멘테이션 모델인 SegFormer를 파인튜닝하는 워크플로우를 상세히 설명했습니다. 이 프로세스는 사전 학습된 가중치와 Hugging Face 생태계를 활용하여, 피자 배달 로봇이 인도와 장애물을 인식하도록 설계된 모델과 같은 특화된 모델을 생성할 수 있게 합니다.

SegFormer 및 시맨틱 세그멘테이션 이해하기

시맨틱 세그멘테이션은 이미지의 모든 개별 픽셀을 분류하는 프로세스로, 표준 이미지 분류보다 더 세밀한 수준의 세부 정보를 제공합니다. 이 기능은 의료 영상 및 자율 주행과 같이 정밀한 경계 감지(예: 인도의 정확한 가장자리 식별)가 필요한 애플리케이션에서 매우 중요합니다.

SegFormer는 2021년 Xie et al.에 의해 소개되었으며, 이전의 합성곱 신경망(CNN) 및 Vision Transformer (ViT) 접근 방식을 개선했습니다. 그 구조는 다음과 같이 구성됩니다:

계층적 트랜스포머 인코더: ViT와 달리 SegFormer는 위치 인코딩(positional encodings)을 사용하지 않습니다.
단순한 MLP 디코더: 인코더의 출력을 처리하는 다층 퍼셉트론 디코더입니다.

데이터 준비 및 증강

효과적인 시맨틱 세그멘테이션을 위해서는 정밀한 세그멘테이션 맵이 포함된 데이터셋이 필요합니다. ADE20k, CityScapes, 또는 BDD100K와 같은 일반적인 데이터셋이 존재하지만, 분포 불일치를 피하기 위해 도메인 특화 데이터가 필요한 경우가 많습니다. 예를 들어, 인도에서 작동하는 로봇은 자동차의 관점이 아닌 인도의 관점에서 캡처된 데이터가 필요합니다.

데이터셋 로딩 및 처리

datasets 라이브러리를 사용하여 커스텀 데이터셋(예: segments/sidewalk-semantic)을 로드하고 훈련 및 테스트 세트로 분할할 수 있습니다. 모델이 올바른 형식으로 데이터를 받도록 하기 위해 SegFormerImageProcessor가 사용됩니다.

온더플라이(On-the-Fly) 변환

디스크 공간과 훈련 속도를 최적화하기 위해 Hugging Face는 set_transform을 통한 변환을 권장합니다. 이는 전체 데이터셋을 미리 전처리하는 대신 데이터를 배치 단위로 즉석에서 준비합니다. 다양한 조명 조건에 대한 모델의 회복력을 높이기 위해, torchvision.transforms.ColorJitter를 훈련 파이프라인에 통합하여 밝기, 대비, 채도, 색조를 무작위로 조정합니다.

파인튜닝 워크플로우

모델 선택

SegFormer는 5가지 모델 크기(B0부터 B5까지)를 제공합니다. 엣지 배포(예: 배달 로봇)를 위해서는 작은 크기(약 14MB)와 효율성 덕분에 B0 모델이 권장됩니다. 파인튜닝 프로세스는 일반적으로 ImageNet-1k에서 사전 학습된 모델(nvidia/mit-b0)로 시작합니다.

훈련 설정

파인튜닝은 Hugging Face Trainer API를 통해 관리됩니다. 주요 설정 파라미터는 다음과 같습니다:

하이퍼파라미터: 학습률(예: 0.00006) 및 에포크 수(예: 50).
평가 지표: 예측된 세그멘테이션 마스크와 정답(ground truth) 사이의 겹침 정도를 측정하기 위해 평균 교집합 대비 합집합(mIoU)이 사용됩니다.
로짓 업스케일링(Logit Upscaling): SegFormer는 원래 이미지 해상도의 1/4 크기(height/4, width/4)로 로짓을 출력하므로, mIoU를 계산하기 전에 레이블 크기에 맞추기 위해 선형 보간법(bilinear interpolation)을 사용하여 로짓을 업스케일링해야 합니다.

추론 및 배포

파인튜닝이 완료되면 모델과 이미지 프로세서를 Hugging Face Hub에 푸시할 수 있습니다. 이를 통해 쉽게 공유하고 호스팅된 추론 API를 통해 실시간 테스트를 위한 추론 위젯을 생성할 수 있습니다.

추론 실행

새로운 이미지에 대해 추론을 수행하려면 다음 단계가 필요합니다:

전처리: SegformerImageProcessor를 사용하여 이미지를 처리합니다.
순방향 패스(Forward Pass): 처리된 이미지를 모델에 통과시켜 로짓을 얻습니다.
재스케일링(Rescaling): nn.functional.interpolate를 사용하여 로짓을 원래 이미지 차원으로 업스케일링합니다.
예측: 클래스 차원에 argmax 연산을 적용하여 최종 픽셀 수준의 카테고리 예측을 수행합니다.