nanoVLM: 순수 PyTorch 로 VLM 훈련하기 가장 간단한 저장소

nanoVLM은 순수 PyTorch 를 사용하여 매우 own Vision Language Model (VLM) 을 훈련하는 데 시작하는 가장 간단한 방법입니다. 무료 티어 Colab 노트북에서 VLM 훈련을 실행할 수 있는 경량 toolkit 입니다.

Andréj Karpathy 의 nanoGPT 에 영감을 받아 시각적 도메인을 위한 유사한 프로젝트를 제공합니다.

nanoVLM 의 핵심은 toolkit입니다. 이 toolkit 는 이미지와 텍스트를 이해하고 이를 기반으로 텍스트를 생성할 수 있는 모델을 구축하고 훈련하는 데 도움을 줍니다. nanoVLM 의 아름다움은 그 간단함에 있습니다. 전체 코드베이스는 의도적으로 최소화되고 가독성이 높아, 초보자나 VLM 의 작동 원리를 이해하고 싶지만 압도당하지 않는 사람들을 위해 완벽한 도구입니다.

이 블로그 포스트에서는 프로젝트의 핵심 아이디어를 다루고 저장소와 상호작용하는 간단한 방법을 제공합니다. 우리는 프로젝트의 세부 사항을 설명할 뿐만 아니라 모든 것을 캡슐화하여 여러분이 빠르게 시작할 수 있도록 합니다.

• Vision Language Model(가상의 언어 모델)은 무엇인가요?
• 저장소 사용법
• 아키텍처
• 자신의 VLM 훈련
• 사전 훈련된 모델에서 추론 실행
• 결론
• 참고 자료

nanoVLM toolkit 를 사용하여 Vision Language Model 을 훈련하는 방법은 다음과 같습니다:

# 저장소 복제
git clone https://github.com/huggingface/nanoVLM.git
# 훈련 스크립트 실행
...

이 Colab 노트북은 로컬 설정 없이 훈련 실행을 시작하는 데 도움이 됩니다!

이름과 같이, Vision Language Model (VLM) 은 시각 및 텍스트 두 모달리티를 처리하는 멀티모달 모델입니다. 이러한 모델은 일반적으로 이미지와/또는 텍스트를 입력으로 받아 텍스트를 출력으로 생성합니다.

이미지와 텍스트의 이해를 바탕으로 텍스트 (출력) 를 생성하는 것은 강력한 패러다임입니다. 이는 이미지 캡셔닝과 객체 감지부터 시각적 콘텐츠에 대한 질문 답변까지 다양한 응용 프로그램을 가능하게 합니다 (아래 표 참조). nanoVLM 이 훈련 목표로서 시각적 질문 답변 (Visual Question Answering) 만 집중한다는 점은 주목해야 할 사항입니다.

| 이미지를 캡션 생성 | 침대에서 두 마리의 고양이가 리모콘 근처에 누워있는 이미지 | 캡셔닝 | |
| 이미지 내 객체 감지 | <locxx><locxx><locxx><locxx> |
객체 감지 | |
| 이미지 내 객체 분할 | <segxx><segxx><segxx> |
의미 분할 | |
| 이미지에는 몇 마리의 고양이가 있나요? | 2 | 시각적 질문 답변 |

VLM 에 대해 더 배우고 싶으시다면, 최신 블로그를 강력히 권장합니다: Vision Language Models (더 잘, 더 빠르며, 더 강력함)

"말하기는 저렴하지만, 코드를 보여주세요" - Linus Torvalds

이 섹션에서는 코드베이스를 안내합니다. 따라가면서 참조용 탭을 열어두는 것이 도움이 됩니다.

아래는 저장소의 폴더 구조입니다. 간결성을 위해 헬퍼 파일을 제거했습니다.

.
├── data
│ ├── collators.py
...

.
├── data
│ └── ...
...

nanoVLM 은 두 잘 알려진 널리 사용된 아키텍처를 모델링합니다. 우리의 시각적 백본 (models/vision_transformer.py)
은 표준 시각적 트랜스포머이며, 구체적으로 Google 의 SigLIP 시각적 인코더입니다. 우리의 언어 백본은 Llama 3 아키텍처를 따릅니다.

시각 (vision) 과 텍스트 (text) 모달리티는 Modality Projection 모듈을 사용하여 align됩니다. 이 모듈은 시각 백본 (vision backbone) 에서 생성된 이미지 임베딩을 입력으로 받아, 언어 모델의 임베딩 레이어에서 생성된 텍스트 임베딩과 호환되는 임베딩으로 변환합니다. 이러한 임베딩은 연결되어 언어 디코더 (language decoder) 에 공급됩니다. Modality Projection 모듈은 픽셀 쉐플 (pixel shuffle) 연산과 선형 레이어 (linear layer) 로 구성됩니다.

픽셀 쉐플은 이미지 토큰의 수를 줄여 계산 비용을 절감하고, 특히 트랜스포머 기반 언어 디코더는 입력 길이에 민감하므로 학습 속도를 높입니다. 아래 그림은 이 개념을 보여줍니다.

모든 파일은 매우 경량화되어 있고 잘 문서화되어 있습니다. 구현 세부 사항을 더 잘 이해하기 위해 개별적으로 확인하는 것을 강력히 권장합니다 (models/xxx.py)

학습 중에는 다음 사전 학습된 백본 (backbone) 가중치를 사용합니다:

• 시각 백본: google/siglip-base-patch16-224
• 언어 백본: HuggingFaceTB/SmolLM2-135M

또한 SigLIP/SigLIP 2 (시각 백본) 와 SmolLM2 (언어 백본) 의 다른 변형으로 백본을 교체할 수 있습니다.

이제 아키텍처에 익숙해졌으니, train.py 를 사용하여 자체 Vision Language Model 을 훈련하는 방법에 대해 넘어가겠습니다.

├── data
│ └── ...
...

훈련을 시작하려면 다음 명령어를 실행할 수 있습니다:

python train.py

이 스크립트는 전체 훈련 파이프라인의 단일 창으로, 다음을 포함합니다:

• 데이터셋 로드 및 전처리
• 모델 초기화
• 최적화 및 로깅

구성 (Configuration)

모든 것 전에, 스크립트는 models/config.py 에서 두 구성 클래스를 로드합니다:

TrainConfig
: 학습에 유용한 구성 파라미터, 예를 들어 학습률, 체크포인트 경로 등.
VLMConfig
: VLM 을 초기화하는 데 사용되는 구성 파라미터, 예를 들어 숨겨진 차원 (hidden dimensions), attention heads 수 등.

데이터 로드 (Data Loading)

데이터 파이프라인의 핵심은 get_dataloaders 함수입니다. 이 함수는:

• Hugging Face 의 load_dataset API 를 통해 데이터셋을 로드합니다.
• 여러 데이터셋 (제공된 경우) 을 결합하고 섞습니다.
• 인덱싱을 통해 train/val 분할을 적용합니다.
• 커스텀 데이터셋 (VQADataset, MMStarDataset) 및 콜라토어 (VQACollator, MMStarCollator) 로 감쌉니다.

유용한 플래그는 data_cutoff_idx 입니다. 작은 부분집합을 디버깅할 때 유용합니다.

모델 초기화 (Model Initialization)

모델은 VisionLanguageModel 클래스를 통해 구축됩니다. 체크포인트에서 복원하는 경우, 매우 간단합니다:

from models.vision_language_model import VisionLanguageModel
model = VisionLanguageModel.from_pretrained(model_path)

그렇지 않으면 시각 및 언어 백본을 선택적으로 미리 로드하여 새로 초기화된 모델을 얻습니다.

최적화기 설정: 두 개의 LR (Optimizer Setup: Two LRs)

모달리티 프로젝터 (MP) 는 새로 초기화되어 있고 백본은 사전 학습되었으므로, 최적화기는 각자 고유한 학습률을 가진 두 개의 파라미터 그룹으로 나뉩니다:

• MP 에 대한 더 높은 LR
• 인코더/디코더 스택에 대한 더 작은 LR

이 균형은 MP 가 빠르게 학습하면서 시각 및 언어 백본의 지식을 보존합니다.

훈련 루프 (Training Loop)

이 부분은 상당히 표준적이지만 신중하게 구조화되어 있습니다:

• 성능을 개선하기 위해 torch.autocast 를 사용하여 미드精度 (mixed precision) 를 사용합니다.

• get_lr 를 통해 코사인 학습률 스케줄링과 선형 워밍업 (linear warmup) 을 구현합니다.

• Token throughput (tokens/sec) 는 성능 모니터링을 위해 배치별로 로그로 기록됩니다.

매 250 단계 (설정 가능)마다 모델은 검증 및 MMStar 테스트 데이터셋에서 평가됩니다. 정확도가 개선되면 모델이 체크포인트가 됩니다.

로그 및 모니터링

만약 log_wandb 가 활성화되어 있다면, batch_loss, val_loss, accuracy, 그리고 tokens_per_second와 같은 훈련 통계는 실시간 추적용 Weights & Biases 로 로그로 기록됩니다.

런은 샘플 크기, 배치 크기, 에포크 수, 학습률 및 날짜와 같은 메타데이터를 사용하여 자동으로 이름을 붙여지며, 이는 보조 함수 get_run_name 으로 처리됩니다.

Hub 로 푸시

다른 사람들이 찾아보고 테스트할 수 있도록 훈련된 모델을 Hub 로 푸시하려면 다음을 사용하세요:

model.save_pretrained(save_path)

다음과 같이 쉽게 푸시할 수 있습니다:

model.push_to_hub("hub/id")

nanoVLM 을 툴킷으로 사용하여 모델을 훈련하고 Hub 로 게시했습니다.
우리는 google/siglip-base-patch16-224 와 HuggingFaceTB/SmolLM2-135M 를 백본으로 사용했습니다. 이 모델은 cauldron 의 약 1.7M 샘플에 대해 단일 H100 GPU 에서 약 6 시간 동안 훈련되었습니다.
이 모델은 SoTA 모델과 경쟁하는 것이 아니라 VLM 의 구성 요소 및 훈련 과정을 해명하기 위한 것입니다.

.
├── data
│ └── ...
...

훈련된 모델을 inference 하려면 generate.py 스크립트를 사용하세요. 다음 명령어로 생성 스크립트를 실행할 수 있습니다:

python generate.py

이것은 기본 인수를 사용하여 이미지 assets/image.png 에 대해 "What is this?" 쿼리를 실행합니다.

본 스크립트를 사용자의 이미지 및 프롬프트와 같이 사용할 수 있습니다:

python generate.py --image path/to/image.png --prompt "You prompt here"

스크립트의 핵심을 시각화하고 싶다면, 다음과 같은 줄들입니다:

model = VisionLanguageModel.from_pretrained(source).to(device)
model.eval()
tokenizer = get_tokenizer(model.cfg.lm_tokenizer)
...

우리는 모델을 생성하고 eval 로 설정합니다. 텍스트 프롬프트를 토큰화하는 토크나이저와 이미지를 처리하는 이미지 프로세서 (image processor) 를 초기화합니다. 다음 단계는 입력을 처리하고 model.generate 를 실행하여 출력 텍스트를 생성하는 것입니다. 마지막으로 batch_decode 를 사용하여 출력을 디코딩합니다.

이미지	프롬프트	생성
What is this?	In the picture I can see the pink color bed sheet. I can see two cats lying on the bed sheet.
What is the woman doing?	Here in the middle she is performing yoga

훈련된 모델을 UI 인터페이스에서 inference 하고 싶다면, Hugging Face Space 를 통해 모델과 상호작용할 수 있습니다.

이 블로그 포스트에서는 VLM 이 무엇인지 설명하고 nanoVLM 을 구동하는 아키텍처 선택을 탐색했으며, 훈련 및 inference 워크플로우를 자세히 설명했습니다.
코드를 경량화하고 가독성을 유지하여 nanoVLM 은 학습 도구이자 구축할 수 있는 기반이 되기를 목표로 합니다. 멀티모달 입력이 어떻게 정렬되는지 이해하고 싶거나, VLM 을 자체 데이터셋에서 훈련하고 싶다면 이 저장소는 시작을 제공합니다.
실험해 보거나, 이를 기반으로 구축하거나, 질문이 있다면 저희는 귀의 목소리를 듣고 싶습니다. Happy tinkering!