[논문 리뷰] SCENIC: A JAX Library for Computer Vision Research and Beyond

Mostafa Dehghani, Alexey A. Gritsenko외 3명|arXiv (Cornell University)|2021. 10. 18.

Advanced Neural Network Applications참고 문헌 10인용 수 5

한 줄 요약

SCENIC은 비전 모델, 특히 트랜스포머 기반 아키텍처의 빠른 프로토타이핑과 대규모 훈련을 위한 오픈소스 JAX 라이브러리입니다. 데이터 파이프라인, 모델 아키텍처, 훈련 루프, 평가를 위한 모듈식이고 재사용 가능한 구성 요소를 제공하여 GPU/TPU 지원을 통해 이미지, 영상 및 다중모odal 작업에서 효율적인 실험을 가능하게 합니다.

ABSTRACT

Scenic is an open-source JAX library with a focus on Transformer-based models for computer vision research and beyond. The goal of this toolkit is to facilitate rapid experimentation, prototyping, and research of new vision architectures and models. Scenic supports a diverse range of vision tasks (e.g., classification, segmentation, detection)and facilitates working on multi-modal problems, along with GPU/TPU support for multi-host, multi-device large-scale training. Scenic also offers optimized implementations of state-of-the-art research models spanning a wide range of modalities. Scenic has been successfully used for numerous projects and published papers and continues serving as the library of choice for quick prototyping and publication of new research ideas.

연구 동기 및 목표

• 새로운 비전 아키텍처를 프로토타이핑하기 위한 통합적이고 확장 가능한 프레임워크를 제공함으로써 컴퓨터 비전 분야의 연구를 가속화하기 위해.
• 트랜스포머 및 MLP 기반 아키텍처를 포함한 비전 모델에 대해 대규모, 다중 장치 및 다중 호스트 훈련을 지원하기 위해.
• 이미지 분류, 세그멘테이션, 검출 및 다중모달 학습과 같은 다양한 작업에서의 빠른 실험을 가능하게 하기 위해.
• ViT, DETR, U-Net와 같은 최신 기술 모델의 최적화되고 검증된 구현을 제공함으로써 엔지니어링 오버헤드를 줄이기 위해.
• 빠른 프로토타이핑과 논문 수준의 연구 코드 모두에 적합한 프로덕션 수준의 코드베이스로 기능하기 위해.

제안 방법

• GPU/TPU에서 자동 미분 및 고성능 계산을 위해 JAX를 활용합니다.
• 함수형 프로그래밍 패턴을 사용하여 모델 정의 및 훈련을 위한 신경망 라이브러리로 Flax를 사용합니다.
• 공유 기능을 위한 모듈식 라이브러리 수준의 구성 요소(예: dataset_lib, model_lib, train_lib)로 코드를 구성합니다.
• 초기 설정부터 전체 아키텍처 재설계에 이르기까지 다양한 작업을 위한 복제 가능하고 최소한의 코드 템플릿을 통해 프로젝트 수준의 커스터마이제이션을 지원합니다.
• 모델 정의를 표준화하기 위해 build_flax_model, loss_fn, get_metrics_fn 구성 요소를 포함한 BaseModel 추상 클래스를 도입합니다.
• 최적화된 트레이너(예: 분류 및 세그멘테이션 트레이너)를 포함한 유연한 훈련 루프 추상화를 사용하며, 이를 확장하거나 대체할 수 있습니다.

실험 결과

연구 질문

• RQ1 통합적이고 모듈식 소프트웨어 프레임워크는 새로운 비전 아키텍처의 개발 및 평가를 어떻게 가속화할 수 있는가?
• RQ2 다양한 장치와 호스트를 통해 비전 모델의 대규모 훈련을 효율적으로 수행하기 위한 설계 원칙은 무엇인가?
• RQ3 최소한의 구성 가능 라이브러리가 유연성을 희생시키지 않고 엔지니어링 복잡성을 얼마나 줄일 수 있는가?
• RQ4 하나의 코드베이스가 이미지, 영상 및 다중모달 학습과 같은 다양한 작업을 얼마나 효과적으로 지원할 수 있는가?
• RQ5 공동 사용 인fra구조는 비전 연구에서 재현 가능성과 코드 재사용을 얼마나 크게 향상시킬 수 있는가?

주요 결과

• SCENIC은 ViViT, OmniNet, TokenLearner, MBT와 같은 여러 최신 기술 모델의 개발 및 발표에 성공적으로 사용되었습니다.
• 최소한의 부록 코드로 빠른 프로토타이핑이 가능하여 연구자들이 인프라에 대한 고민보다 모델 혁신에 집중할 수 있습니다.
• TPU와 GPU에서의 다중 호스트, 다중 장치 훈련 지원으로 대규모 비전 연구에 대한 확장성을 확보합니다.
• 모듈식 설계 덕분에 사전 구축된 모델과 트레이너를 사용하는 것에서부터 완전한 커스터마이제이션까지 다양한 프로젝트를 구현할 수 있으며, 강제적인 추상화 오버헤드가 없습니다.
• 라이브러리 수준의 구성 요소는 최소화되고 잘 테스트되어 있으며, 유저가 널리 재사용할 수 있는 功能만 업스트림으로 이관되어 있어 명확성과 유지보수성을 유지합니다.
• 이 프레임워크는 수많은 구글 연구 프로젝트에서 주요 코드베이스로 채택되어, 프로덕션 연구 환경에서의 견고성과 확장성의 우수성을 입증했습니다.