[論文レビュー] SCENIC: A JAX Library for Computer Vision Research and Beyond

Mostafa Dehghani, Alexey A. Gritsenko+3 more|arXiv (Cornell University)|Oct 18, 2021

Advanced Neural Network Applications参考文献 10被引用数 5

ひとこと要約

SCENIC は、特にトランスフォーマーに基づくアーキテクチャを対象として、ビジョンモデルの迅速なプロトタイピングおよび大規模学習を目的としたオープンソースの JAX ライブラリです。データパイプライン、モデルアーキテクチャ、トレーニングループ、評価のためのモジュラーで再利用可能なコンponentsを提供し、GPU/TPU 対応により、画像、動画、マルチモーダルタスクにおいても効率的な実験を可能にします。

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）にコードを整理すること。
• ハイパーパramータチューニングからアーキテクチャの完全な再設計まで、さまざまなタスクに対応できるフォーク可能で最小限のコードのテンプレートをプロジェクトレベルでカスタマイズ可能にすること。
• モデル定義を標準化するため、BaseModel 抽象クラスを導入し、build_flax_model、loss_fn、get_metrics_fn のコンponentsを統一すること。
• 最適化されたトレーナー（例：分類用・セグメンテーション用トレーナー）を備えた柔軟なトレーニングループの抽象化を採用し、拡張や置き換えが可能になるようにすること。

実験結果

リサーチクエスチョン

• RQ1 統合的でモジュラーなソフトウェアフレームワークは、新しいビジョンアーキテクチャの開発と評価をどのように加速できるか？
• RQ2 複数のデバイスおよびホストをまたがるビジョンモデルの大規模学習を効率的に行うために、どのような設計原則が必要か？
• RQ3 最小限で組み合わせ可能なライブラリは、柔軟性を損なわずに、エンジニアリングの複雑さをどの程度軽減できるか？
• RQ4 1つのコードベースが、画像、動画、マルチモーダルラーニングを含む多様なタスクを効果的にサポートできるか？
• RQ5 共通のインfraストラクチャは、ビジョン研究における再現性とコード再利用性を顕著に向上させられるか？

主な発見

• SCENIC は、ViViT や OmniNet、TokenLearner、MBT といった複数の最先端モデルの開発および公開に成功して使用されています。
• 最小限のボイラープレートで迅速なプロトタイピングが可能であり、研究者がインfraストラクチャに費やす時間ではなく、モデルのイノベーションに集中できるようになっています。
• TPU や GPU 上でのマルチホスト・マルチデバイストレーニングをサポートしており、大規模ビジョン研究におけるスケーラビリティを確保しています。
• モジュラー設計により、事前構築済みモデルやトレーナーの使用から完全なカスタマイズまで、さまざまなプロジェクトのニーズに対応でき、強制的な抽象化のオーバーヘッドが生じません。
• ライブラリレベルのコンponents は最小限に抑えられ、広く再利用可能な機能のみをアップストリームに統合しており、コードの明確さと保守性を保っています。
• フレームワークは、多数の Google 研究プロジェクトの主要コードベースとして採用されており、生産的で拡張性のある研究環境においてその堅牢性と実用性が裏付けられています。