[论文解读] 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 的开源库，旨在加速视觉模型（尤其是基于 Transformer 的架构）的快速原型设计与大规模训练。它提供了模块化、可重用的组件，涵盖数据管道、模型架构、训练循环和评估，支持在 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.

研究动机与目标

• 通过提供一个统一、可扩展的框架，加速新型视觉架构的原型设计，推动计算机视觉研究的发展。
• 支持视觉模型（包括 Transformer 和 MLP 基础架构）在大规模、多设备及多主机环境下的训练。
• 支持在图像分类、分割、检测及多模态学习等多样化任务中实现快速实验。
• 通过提供经过优化、充分测试的最先进模型（如 ViT、DETR 和 U-Net）实现，降低工程复杂度。
• 作为生产级代码库，既适用于快速原型设计，也适用于发布级别的研究代码。

提出的方法

• 利用 JAX 实现自动微分与在 GPU/TPU 上的高性能计算。
• 使用 Flax 作为神经网络库，通过函数式编程模式定义和训练模型。
• 将代码组织为模块化的库级组件（如 dataset_lib、model_lib、train_lib），以实现功能共享。
• 通过可 fork、代码量极少的项目级模板，支持从超参数调优到完整架构重构等各类自定义需求。
• 引入 BaseModel 抽象类，通过 build_flax_model、loss_fn 和 get_metrics_fn 组件统一模型定义方式。
• 采用灵活的训练循环抽象，配备优化的训练器（如分类和分割训练器），支持扩展或替换。

实验结果

研究问题

• RQ1 统一、模块化的软件框架在多大程度上能加速新型视觉架构的开发与评估？
• RQ2 哪些设计原则能够实现跨多设备和多主机的视觉模型高效大规模训练？
• RQ3 一个极简、可组合的库在不牺牲灵活性的前提下，能在多大程度上降低工程复杂度？
• RQ4 单一代码库在多大程度上能有效支持图像、视频及多模态学习等多样化任务？
• RQ5 共享基础设施在多大程度上能显著提升视觉研究中的可复现性与代码复用性？

主要发现

• SCENIC 已成功用于开发并发表多个最先进模型，包括 ViViT、OmniNet、TokenLearner 和 MBT。
• 该库通过极简的样板代码实现快速原型设计，使研究人员能专注于模型创新而非基础设施搭建。
• 对 TPU 和 GPU 上的多主机、多设备训练的支持，确保了大规模视觉研究的可扩展性。
• 模块化设计使项目可灵活选择使用预构建模型与训练器，或实现完全自定义，且无强制抽象开销。
• 库级组件保持极简且经过充分测试，仅将广泛复用的功能上游化，从而保障代码清晰性与可维护性。
• 该框架已被多个 Google 研究项目采纳为主要代码库，证明其在生产级研究环境中的稳健性与可扩展性。