Upstage AI Lab 5기

한식 이미지 생성 모델

Shining Future 2025. 4. 26. 13:20

Upstage AI Lab 5기 교육과정 중 마지막 프로젝트로 기업 연계 프로젝트를 진행하였다.

  • 기업명: 커넥트브릭(connectbrick)
  • 팀 주제: 한국 음식 메뉴 이미지 생성 모델(Text-to-Image) 개발
  • 과제 추진 배경: 대부분의 오픈소스 모델들은 한국 음식에 대한 이해도가 낮아, 실제와 거리가 있는 이미지를 생성
  • 목표: 고퀄리티 실사 이미지 기반의 한국 음식 이미지 생성 모델을 학습 및 파인튜닝하여, 보다 실제 사진 같은 이미지를 생성
  • 주요 작업
    • 고품질 한식 이미지 수집 및 캡셔닝
    • 다양한 T2I 모델(SD, Flux, Janus 등) 파인튜닝 시도
    • LoRA 기반 경량 모델 학습 실험
    • ControlNet, VAE 등 보조기법과의 통합 실험
  • 기간: 2025. 3. 21 ~ 2025. 4. 14

구체적인 이미지 생성 모델 개발 과정은 아래와 같이 진행하였다.


  원래 계획은 위와 같이 K-Food 메뉴를 한 가지씩 계속 추가하면서 모델을 학습시켜 다양한 메뉴를 생성할 수 있는 모델을 개발하는 것이었느나 실제 과제를 진행하는 과정에서 당면한 가장 큰 issue는 주어진 computing resource (RAM, GPU, VRAM, Disk, ...)의 한계로 인해 발생하는 OOM(Out of memory) error와 disk 사용 한도 초과에 따른 GPU server down 문제를 해결하는 것이었다. 

 

  disk 사용 한도 초과에 따른 GPU server down 문제를 해결하기 위해서 프로그램 수행하면서 disk 사용량을 실시간 모니터링하여 위험 수준에 도달할 경우 프로그램이 자동 종료되게 하여 GPU server down은 발생하지 않도록 하였고 이미지나 체크포인트 모델과 같은 대용량 파일들은 GCS(Google Cloud Server)에 저장해 두고 필요할 때만 GPU server의 disk에 읽어 들였다가 사용 후에는 disk에서 삭제하는 방법을 사용하였다.

 

  OOM error issue는 아래와 같이 다양한 방법을 동시에 적용하여 해결하였다. 

✅ OOM 방지를 위한 실제 적용 전략 (우선순위별 정리)

  1. Adafactor 옵티마이저 사용
    • 가장 큰 효과를 얻은 전략이었다. 일반적으로 사용하는 Adam이나 AdamW 대신 Adafactor를 선택함으로써, 추가적인 옵티마이저 상태 저장 메모리를 대폭 줄였다. (Adafactor optimizer는 특히 SDXL처럼 파라미터가 수억 개 이상인 대형 모델을 학습할 때 OOM 발생 가능성을 크게 낮출 수 있게 함.)
    • 실제 코드에서는 scale_parameter=True, relative_step=True 등의 설정을 통해 학습률 스케줄링까지 포함시켰다.
  2. bfloat16 데이터 타입 사용
    • 파라미터와 연산 결과를 float32 대신 bfloat16으로 처리하여 약 2배의 메모리 절약을 실현하였다.
    • bfloat16은 메모리 사용량은 줄이되, 수치적 안정성은 float16보다 더 뛰어난 포맷이라 OOM 방지와 성능 사이의 균형에 적합
    • 모델 로딩 및 학습 모두에 이 설정을 반영하였다.
  3. gradient_accumulation_steps 설정
    • 한 번에 큰 배치를 처리하지 않고, 작은 배치를 여러 번 처리하여 gradient를 누적하는 방식으로 구현하였다.
    • 이를 통해 실질적인 배치 크기는 유지하면서도 GPU 메모리 요구량은 줄일 수 있었다.
    • 실제 코드 적용 조건: batch_size=4에 gradient_accumulation_steps=4로 설정하여 메모리 부담을 낮추면서 16의 배치 크기로 학습
  4. gradient checkpointing 활성화
    • 모델의 중간 레이어 출력을 저장하지 않고 필요할 때 재계산함으로써 메모리를 절약
    • 특히 SDXL 같이 레이어가 깊고 복잡한 모델에 매우 효과적이며, 코드상에서는 pipeline.unet.enable_gradient_checkpointing()을 통해 적용
  5. xformers 기반 memory-efficient attention 사용
    • attention 연산에서 메모리를 절약할 수 있는 xformers 백엔드를 사용하도록 설정
    • 이를 통해 multi-head attention 과정의 메모리 사용량을 크게 줄였다.
    • 코드에서는 set_attn_processor() 및 enable_xformers_memory_efficient_attention() 함수들을 통해 적용
  6. VAE 최적화 기능 사용 (slicing 및 tiling)
    • VAE 인코더와 디코더 처리 시 전체 이미지를 한 번에 처리하지 않고 조각내어(slicing/tiling) 처리하여 메모리 부담을 줄였다.
    • enable_vae_slicing() 및 enable_vae_tiling()을 통해 적용
  7. 주기적인 메모리 정리
    • 학습 중 50 step마다 gc.collect()와 torch.cuda.empty_cache()를 호출하여 불필요한 메모리를 강제로 정리하였다.
  8. 실시간 자원 모니터링 및 대응 시스템
    • DRAM, VRAM, 디스크 사용량을 주기적으로 확인하고, 임계치를 초과하면 자동으로 캐시 정리 또는 학습 종료 등의 조치를 취하도록 구성되어 있는 resource_monitor 모듈을 작성하였고 학습을 실행하는 train.py에서 별도 스레드로 실행

  위와 같이 다양한 시도를 한 끝에 OOM error issue와 disk 사용 한도 초과 문제를 해결하기는 하였으나 아쉽게도 본격적인 모델 학습과 검증/테스트는 얼마 진행하지 못한 채로 과제 수행 기간이 종료되었다. 

  과제 수행 기간은 종료되었지만 개인적인 추가 학습 차원에서 과제 종료 후에도 추가 학습/검증/테스트를 진행해 보려고 한다.