Transformers.js를 사용한 기계 학습 기반 웹 게임 만들기

Hugging Face는 사용자 브라우저 내에서 완전히 실행되는 실시간 기계 학습(ML) 기반 웹 게임인 Doodle Dash의 개발을 상세히 설명했습니다. Transformers.js를 활용하여 서버 지연을 제거함으로써, 사용자가 그릴 때 모델이 초당 60회 이상의 예측을 수행할 수 있게 됩니다.

모델 훈련 및 아키텍처

Doodle Dash의 핵심은 Google의 "Quick, Draw!" 데이터셋의 하위 집합에서 훈련된 스케치 감지 모델로, 345개의 카테고리에 걸쳐 5백만 개 이상의 drawings가 포함되어 있습니다.

모델 선택

개발자는 ImageNet-1k에서 사전 훈련된 경량 Vision Transformer(ViT)인 apple/mobilevit-small을 파인튜닝했습니다. 이 모델은 모바일 친화적인 아키텍처와 작은 footprint 때문에 선택되었으며, 560만 개의 파라미터와 약 20 MB의 파일 크기로 구성되어 브라우저 내 실행에 적합합니다.

파인튜닝 프로세스

파인튜닝 워크플로는 다음과 같은 기술적 단계로 구성되었습니다:

  1. 데이터 로드: "Quick, Draw!" 데이터셋 하위 집합 가져오기.
  2. 전처리: MobileViTImageProcessor를 사용하여 데이터 변환.
  3. 구성: 콜레이트 함수와 평가 지표 정의.
  4. 모델 초기화: 사전 훈련된 MobileViTForImageClassification 모델 로드.
  5. 훈련: TrainerTrainingArguments 헬퍼 클래스 활용.
  6. 평가: 🤗 Evaluate 라이브러리를 사용하여 성능 검증.

Transformers.js를 사용한 브라우저 배포

Transformers.js는 백엔드 서버 없이 브라우저에서 직접 🤗 Transformers 모델을 실행할 수 있게 해주는 JavaScript 라이브러리입니다. Python 라이브러리와 기능적으로 동등한 API를 제공합니다.

ONNX로 모델 변환

Transformers.js는 실행에 ONNX Runtime을 사용하기 때문에 PyTorch 모델은 ONNX 형식으로 변환해야 합니다. 이는 Transformers.js 저장소에 제공된 변환 스크립트를 통해 🤗 Optimum 라이브러리를 사용하여 수행됩니다.

python -m scripts.convert --model_id <model_id>

기술적 구현

계산량이 많은 추론 프로세스가 메인 UI 스레드(렌더링 및 사용자 입력을 처리)를 차단하지 못하도록 방지하기 위해, 게임은 Web Workers API를 구현합니다. 추론 로직은 별도의 워커 스레드(worker.js)에 격리되어 있으며, 여기서 image-classification 파이프라인이 초기화되고 회색조 이미지가 처리되어 예측을 반환합니다.

게임 디자인 및 최적화

실시간 추론 루프

몇 초마다 예측을 수행하는 원래의 "Quick, Draw!" 게임과 달리, Doodle Dash는 브라우저 내 추론의 고주파 성능을 활용하여 더 빠른 게임 루프를 만듭니다:

  • 목표: 플레이어는 60초 이내에 가능한 많은 낙서를 그리려고 합니다.
  • 메커니즘: 올바른 예측이 이루어지면 캔버스가 즉시 지워져 새로운 단어가 제시됩니다.
  • 페널티: 단어를 건너뛰면 남은 시간에서 3초가 차감됩니다.
  • 점수 조정: 모델이 단순히 레이블을 지우는 것을 방지하기 위해, 게임은 처음 n개의 잘못된 레이블에 대해 점수를 감소시키며, n은 시간이 지남에 따라 증가합니다.

데이터셋 정제

모델 크기가 작음(~20MB)에도 불구하고 게임 품질을 유지하기 위해 개발자는 원래의 345개 클래스를 필터링했습니다. 레이블은 다음 조건에 해당하는 경우 제거되었습니다:

  • 다른 레이블과 너무 유사한 경우 (e.g., "barn" vs. "house)).
  • 충분히 자세히 이해하거나 그리기 어려운 경우 (e.g., "animal migration" 또는 "brain)).
  • 애매한 경우 (e.g., "bat)).

이 필터링 과정을 통해 300개 이상의 서로 다른 클래스로 구성된 최종 세트가 얻어졌습니다.

Sources