SCENIC：用于计算机视觉研究及其他领域的JAX库

SC E N I C：用于计算机视觉研究及其他领域的JAX库谷歌大脑dehghani@google.comAlexey Gritsenko谷歌大脑agritsenko@google.com谷歌研究aarnab@google.comMatthiasMinderer谷歌大脑mjlm@google.com谷歌研究yitay@google.com摘要S CENIC 是 一 个 开 源 的 1JAX 库 ， 专 注 于 基 于transformer的模型，用于计算机视觉研究等。该工具包的目标是促进新体系结构和模型的快速实验、原型设计和研究。S CENIC支持各种任务（例如，分类，分割 ， 检 测 ） ， 并 促 进 多 模 式 问 题 的 工 作 ， 以 及GPU/TPU支持大规模，多主机和多设备训练。S CENIC还提供最先进的研究模型的优化实施，涵盖广泛的模型。S CENIC已成功用于许多项目和发表的论文，并继续作为快速原型和发表新研究想法的首选图书馆。1. 介绍这是一个令人兴奋的时刻，在计算机视觉研究使用基于注意力的模型。随着像ViT [12]这样的新架构席卷全球，对软件和机器学习基础设施的需求明显，以支持简单和可扩展的神经网络架构研究。随着注意力模型[6，12，19，28]和仅MLP [30]架构变得越来越流行，我们预计在未来几年将看到更多的研究推动该领域向前发展。我们介绍了SCENIC，这是一个开源的JAX库，用于视觉及其他领域的快速和可扩展的研究。SCENIC已成功用于开发图像、视频和音频等模态（包括多模态设置）的分类、分割和检测模型1https://github.com/google-research/scenicSCENIC通过提供上游预训练所需的管道，为迁移学习提供了一种高效且易于使用的设置使用大规模数据和/或自我监督训练目标，以及下游评估，例如，零次学习、少次学习、线性探测以及对多个任务和数据集的全面微调。SCENIC致力于成为一个统一的，一体化的代码库，以满足建模需求，目前在各种任务中提供最先进的模型实现，如ViT [12]，DETR [8]，CLIP [22]，MLPMixer [30]，T5 [23]，BERT [11][13]和U-Net [24]。最重要的是，SCENIC已被用于许多Google项目和研究论文，如ViViT [6]，OmniNet [28]，TokenLearner [25]，MBT [21]、PolyViT [19]、MTV [31]、缩放研究行为[5，10，29]和效率[9]的各种模型等[15我们预计在不久的将来，会有更多不同风格的研究项目在SCENIC存储库中开源。SCENIC是在JAX [7]中开发的，使用Flax [14]作为神经网络库，依赖于TFDS [4]和DMVR [2]来实现大多数任务和数据集的输入管道，并利用CLU [1]提供的常见训练循环功能。JAX是一个超简单易用的库，可以自动区分原生Python和NumPy函数。此外，它支持在加速器（包括GPU和TPU）上进行多主机和多设备训练，使其成为大规模机器学习研究的理想选择。简而言之，SCENIC是（i）一组共享的轻量级库，用于解决在视觉和其他方面训练大规模（即多设备，多主机）模型时经常遇到的任务;以及（ii）一些项目，包含使用这些库完全充实的特定问题训练和评估循环。SCENIC旨在提出不同层次的AB-大多数项目使用的通用代码。自定义基线代码。定制项目代码。图1. S CENIC中的代码被组织在项目级部分中，这是针对特定项目或基线或库级部分的定制代码，它实现了大多数项目所适应的通用功能和通用模式。牵引它支持从那些只需要改变超参数的项目，到那些需要在输入管道、模型架构、损失和指标以及训练循环上进行定制的项目。为了实现这一点，SCENIC中的代码被组织为项目级代码（指特定项目或基线的定制代码）或库级代码（指大多数项目采用的通用功能和通用模式）。1）。项目级代码位于projects目录中SCENIC旨在促进大规模模型的快速原型化。为了使代码易于理解和扩展，SCENIC设计更喜欢分叉和复制粘贴，而不是增加复杂性或增加抽象。我们只在功能被证明在多个模型和任务中广泛有用在库级代码中最小化对各种用例的支持有助于我们避免积累导致代码复杂和难以理解的泛化。请注意，任何级别的复杂性或抽象都可以添加到项目级代码中。2. JAX和Flax在介绍SCENIC的设计之前，我们先简要介绍一下JAX的背景知识[7]。JAX提供了Python和NumPy程序的可组合转换，这些转换提供了差异化（例如，梯度）、矢量化（例如，vmap）、JIT编译到GPU/TPU（例如，jit），以及更多。JAX被设计为纯静态操作（PSC）功能。这意味着JAX转换仅适用于以下函数：（1）不使用全局状态，（2）确定性，以及（3）可表示为一组参数上的静态数据依赖图加速线性代数（XLA）运算符）。许多机器学习系统可以实现为一组PSC Python函数，并使用JAX转换，我们可以很容易地利用XLA的高效率来运行这些系统的计算繁重部分。虽然JAX以函数转换的形式提供了开发机器学习系统所需的构建块，但在它之上还有其他库来满足研究人员的更高层次需求。最值得注意的是Flax [14]，它是一个用于JAX的神经网络库，提供神经网络层以及实用程序和模式。SCENIC使用JAX和Flax来开发一个高效的框架，用于敏捷开发研究思想，在大规模，但灵活的设置。3. SCENIC设计SCENIC提供了一个统一的框架，它足够灵活，可以支持各种需求的项目，而无需编写复杂的代码。SCENIC包含一组研究模型的优化实现，这些模型在广泛的模式（视频，图像，音频和文本）上运行，并支持多个数据集。这也是由于其灵活和低开销的设计。在本节中，我们将介绍不同的部分，并讨论用于组织项目和库代码的结构。项目编号N…基本模型模型库层巴瑟利讷项目#1数据集库列车库公用库项目Project #2场景模型模板/基类。还有效用函数、损失函数和度量函数数据集的输入管道以及一些用于处理原始数据公共层（亚麻模块），重点关注基于注意力的应用程序优化的培训和评估培训循环（使用或分叉）以及一些常见的培训实用程序重新实施一套基准模型常见的实用程序，例如，用于调试你那了不起的项目3.1. 库级代码目标是保持库级代码最小化并经过良好测试，避免引入额外的抽象来支持次要用例。SCENIC提供的共享库分为：• datasetlib：实现IO管道，用于加载和预处理常见任务和基准的数据。所有管道都被设计为可扩展的，并支持多主机和多设备设置，负责在多个主机之间划分数据，不完整的批处理，缓存，预取等。• 模型库：提供若干抽象模型接口（例如，模型库/基础模型中的ClassificationModel或SegmentationModel），具有特定于任务的损失和度量;模型库/层中的神经网络层，专注于注意力和Transformer原语的有效实现;最后，模型库/匹配器中的二分匹配算法[ 3 ]的加速器友好实现。• trainlib：提供用于构建训练循环的工具，并实现了几个优化的训练器（例如，分类训练器和分段训练器），其可以分叉以用于定制。• common lib：通用工具，如日志和调试模块，以及处理原始数据的功能。3.2. 项目级代码SCENIC支持通过“项目”概念为专门任务和数据开发定制解决方案。对于一个项目应该重用多少代码，没有一个放之四海而皆准的方法项目可以仅由配置文件组成，并使用位于库级代码中的公共模型、训练器、任务/数据，或者它们可以简单地派生任何提到的功能并重新定义层、损失、指标、日志记录方法、任务、架构以及训练和评估循环。库级代码的模块化使其足够灵活，可以支持从“按原样运行”到“完全定制”的任何项目常见基线，如Vision Transformer（ViT）、DETR、CLIP、T5、BERT、MLP-Mixer、ResNet、UNet等，在项目/基准中执行。在这个目录中创建分支模型对于新项目来说是一个很好的起点。3.3. SCENIC基础模型解决方案通常包括几个部分：数据/任务管道、模型架构、损失和度量、训练和评估等。为了尝试不同的体系结构，SCENIC引入了“模型”的概念，以便于“插入/拔出”实验。SCENIC模型被定义为网络体系结构、用于在训练期间更新网络权重的损失以及 用 于 评 估 网 络 输 出 的 度 量 。 它 被 实 现 为BaseModel。BaseModel是一个有三个成员的抽象类：构建flax模型、损失fn和获取度量fn。buildflaxmodel函数返回flax模型。一典型的使用模式描述如下：#获取模型类：model_cls = model_lib.models.get_model_cls（“fully_connected_classification”）#构建模型、指标和损失模型= model_cls（config，dataset.Meta_data）#初始化模型参数flax_model =model.build_flax_model dummy_input= jnp.zeros（input_shape，model_input_dtype）model_state，params = flax_model.init（rng，dummy_input，train=False）.pop（“params”）这是如何调用模型：变量={#可训练参数“params”：params，#模型状态#（例如，来自BatchNorm的批次统计数据）**model_state}logits，new_model_state = flax_model.apply（variables，inputs，.）摘要类为限定SCENIC模型在模型库/基本模型中声明。这些模型包括所有模型继承的基本模型， 以及分别定义分类、序列到序列和分段任务的损失和度量的ClassificationModel、MultiLabelClassificationModel、EncoderDecoderModel和SegmentationModel。根据其需求，SCENIC项目可以定义新的基类或覆盖其特定任务，损失和指标的现有基类。SCENIC中的典型模型损失函数期望预测和一批数据：#损失函数：loss_fn（logits：jnp.ndarray，batch：Dict[str，jnp.ndarray]）->float最后，一个典型的get metrics fn返回一个可调用的metric fn，它计算适当的metrics并将它们作为Python字典返回。对于每个度量，度量函数在小批量上计算f（xi，yi），其中xi和yi是第i个示例的输入和标签，并将从度量名称到度量值的元组和度量规范化器（通常是小批量中的示例的数量）的字典它有API：#度量函数：01 - 02 - 2016张国荣（logits：jnp.ndarry，label：jnp.ndarry，）-> Dict[str，Tuple[float，int]]给定从所有主机中的所有设备处理的示例中收集的度量和规范化器值，然后模型训练器负责聚合和计算完整的端到端上游训练和下游微调/评估代码，支持多主机/多设备上的数据并行性，可用于所有SCENIC基线。请注意，对于这些基线模型中的每一个，在尺寸和配置方面，不同变体的检查点都是可用的。SCENIC包含许多最先进的模型，以及训练/评估代码以及用于分类和图像模态之外的任务的检查点。表2提供了几个例子，显示任务和方式的多样性。这些模型要么最初在SCENIC中被定义，要么成为它们所解决的任务的事实上的基线（例如，ViViT），或者小心地移植到SCENIC，不仅可以重现原始结果，而且可以在使用JAX的不同加速器上高效运行。这已经使许多研究人员能够在现有SOTA模型的基础上进行构建，并在不同的设置和模式上尝试新的想法，这是SCENIC的主要任务之一。表2. SCENIC提供的其他型号。完整的训练和评估代码再现了原始结果，所有这些模型都是公开的。模型模式任务美国[8]BERT [11]文本自监督预训练，评估示例的标准化度量值美国[6]GLUE/SuperGLUE基准测试视频分类重要的是，虽然上面的设计模式是推荐的，并且已经被发现在一系列项目中工作得很好，但它不是强制的，并且在项目中不存在偏离上述结构的问题4. 应用SCENIC已用于各种应用和任务。SCENIC不仅提供了不同的指标来评估不同方法在手头任务上的质量，如分类的准确性，而且还提供了统一的工具来评估方法在训练和推理吞吐量，参数大小和FLOP方面的效率。在本节中，我们将展示一些帝国的结果，展示SCENIC中的一些基线。表1.在SCENIC中实现的图像分类模型。TokenLearner [25] 图像，视频图像/视频分类OmniNet[28]图像图像分类[19]第19话图像/音频/视频分类对于上述每种方法以及SCENIC中实现的其他方法，SCENIC的GitHub存储库中有一个专用的README，包含有关详细配置、实验结果、如何训练、如何评估的信息，以及方法和参考文献的简要说明。5. 结论机器学习（ML）基础设施是ML研究的基石。使研究人员能够快速尝试新的想法，并在出现新想法时迅速扩大规模承诺，加速研究。 此外，历史表明，认为利用当时可用的计算的方法通常是最有效的[27]。SCENIC体现了我们开发最佳研究基础设施的经验，我们很高兴与更广泛的社区分享。我们希望看到使用SCENIC开发出更多出色的想法，为实现令人惊叹的表1给出了在SCENIC中实现的图像分类任务的示例基线模型。的ML社区通过AI改善生活所取得的进展。模型预训练数据集ImageNet Top-1GFLOPS参数（M）ResNet-50 [1]-77.083.9725.5MLP-混合器-B/16[30]-76.1412.759.8[26]第二十六话-79.7317.686.5MLP-混合器-L/16ImageNet21K80.2312.759.8确认SCENIC在许多令人惊叹的贡献者的帮助下得到了开发和改进。我们要特别感谢Dirk Weissenborn、AndreasSteiner、Marvin Ritter、Aravindh Mahendran、SamiraAbnar 、 Sunayana Rane 、 Josip Djolonga 、 LucasBeyer 、 Alexander Kolesnikov 、 Xiaohua Zhai 、 RobRomijnders、Rianne van den Berg、JonathanHeek、OlivierTeboul、MarcoCuturi、LuJiang、MarioLucic' 和NeilHoulsby对SCENIC的直接或间接贡献。引用[1] CLU-commonlooputils。http://github.com/google/CommonLoopUtils.1[2] DMVR：DeepMind视频阅读器。https：//github.com/deepmind/dmvr. 1[3] 最佳运输工具（OTT），一个工具箱，为每件事沃瑟斯坦。https://github.com/google-research/ott. 3[4] TensorFlow Datasets，即用型数据集的集合。https://www.tensorflow.org/datasets网站。1[5] Samira Abnar，Mostafa Dehghani，Behnam Neyshabur，and Hanie Sedghi.探索大规模预培训的极限。arXiv预印本arXiv：2110.02095，2021。1[6] Anurag Arnab ， Mostafa Dehghani ， Georg Heigold ，Chen Sun，MarioLu c. ViV iT：视频视觉Transformer。arXiv预印本arXiv：2103.15691，2021。1、4[7] 詹姆斯·布拉德伯里冰霜斯蒂格彼得Hawkins、Matthew James Johnson、Chris Leary、Dougal Maclau-rin、George Necula、Adam Paszke、Jake VanderPlas、SkyeWanderman-Milne 和 QiaoZhang 。 JAX ：Python+NumPy程序的可组合转换，2018。一、二[8] 尼古拉斯·卡里昂、弗朗西斯科·马萨、加布里埃尔·辛纳夫、尼古拉斯·乌斯尼尔、亚历山大·基里洛夫和谢尔盖·扎戈鲁伊科。使用变压器进行端到端对象检测在EuropeanConferenceon ComputerVision 中 ， 第213Springer，2020年。1、4[9] Mostafa Dehghani、Anurag Arnab、Lucas Beyer、AshishVaswani和Yi Tay。效率用词不当。arXiv预印本arXiv：2110.12894，2021。1[10] Mostafa Dehghani，Yi Tay，Alexey A Gritsenko，ZheZhao，Neil Houlsby，Fernando Diaz，Donald Metzler，and Oriol Vinyals. 基 准 彩 票 。 arXiv 预 印 本 arXiv ：2107.07002，2021。1[11] 雅各布·德夫林张明伟李肯顿和克里斯蒂娜·图塔诺娃。Bert：为语言理解而进行的深度双向转换器的预训练。arXiv预印本arXiv：1810.04805，2018。1、4[12] AlexeyDosovitskiy，LucasBeyer，AlexanderKolesnikov，Dirk Weissenborn，Xiaohua Zhai，ThomasUnterthiner ， Mostafa Dehghani ， Matthias Minderer ，Georg Heigold，Syl- vain Gelly，et al. An image is worth16x16 words ： Trans- formers for image recognition atscale. arXiv预印本arXiv：2010.11929，2020。1[13] Kaiming He，Xiangyu Zhang，Shaoying Ren，and JianSun.用于图像识别的深度残差学习。在CVPR，2016年。1、4[14] Jonathan Heek ， Anselm Levskaya ， Avital Oliver ，Marvin Ritter，Bertrand Rondepierre，Andreas Steiner，and Marc van Zee.Flax ： A neural network library andecosystem for JAX，2020.一、二[15] Charles Herrmann ， Kyle Sargent ， Lu Jiang ， RaminZabih，Huiwen Chang，Ce Liu ，Dilip Krishnan，andDeqing Sun.金字塔对抗训练提高了vit的表现。arXiv预印本arXiv：2111.15121，2021。1[16] Ajay Jain，Ben Mildenhall，Jonathan T Barron，PieterAbbeel，and Ben Poole.零射击文本引导对象生成与梦想领域。arXiv预印本arXiv：2112.01455，2021。1[17] Gang Li，Gilles Baechler，Manuel Tragut，and Yang Li.学习去噪原始移动用户界面布局，以改善大规模数据集。arXiv预印本arXiv：2201.04100，2022。1[18] Yang Li，Gang Li，Xin Zhou，Mostafa Dehghani，andAlexey Gritsenko. Vut：用于多模态多任务用户界面建模 的 多 功 能 ui Transformer 。 arXiv 预 印 本 arXiv ：2112.05692，2021。1[19] 瓦列里·利霍舍斯托夫、阿努拉格·阿纳布、克日什托夫·乔罗曼斯基、马里奥·卢西奇、易泰、阿德里安·韦勒和穆斯塔法·德赫加尼。Polyvit：在图像、视频和音频上 共 同 训 练 视 觉 转 换 器 。 arXiv 预 印 本 arXiv ：2111.12993，2021。1、4[20] Chengzhi Mao ， Lu Jiang ， Mostafa Dehghani ， CarlVondrick，Rahul Sukthankar，and Irfan Essa.离散表示增强 了 视 觉 Transformer 鲁 棒 性 。 arXiv 预 印 本 arXiv ：2111.10493，2021。1[21] Arsha Nagrani ， Shan Yang ， Anurag Arnab ， ArenJansen，Cordelia Schmid，and Chen Sun.多模态融合的注意力瓶颈。arXiv预印本arXiv：2107.00135，2021。1[22] Alec Radford，Jong Wook Kim，Chris Hallacy，AdityaRamesh ， Gabriel Goh ， Sandhini Agarwal ， GirishSastry，Amanda Askell，Pamela Mishkin，Jack Clark，et al. Learn- ing transferable visual models from naturallanguage super- vision. 国际机器学习会议，第8748-8763页。PMLR，2021年。1[23] Colin Raffel、Noam Shazeer、Adam Roberts、KatherineLee、Sharan Narang、Michael Matena、Yanqi Zhou、Wei Li和Peter J Liu。探索迁移学习的限制与统一的文本 到 文 本 Transformer 。 arXiv 预 印 本 arXiv ：1910.10683，2019。1[24] Olaf Ronneberger，Philipp Fischer，and Thomas Brox. U-net：用于生物医学图像分割的卷积网络。在2015年医学图像计算和计算机辅助干预国际会议上。1[25] Michael S Ryoo ， AJ Piergiovanni ， Anurag Arnab ，MostafaDehghani，andAneliaAngelova.TokenLearner：8个学习过的令牌可以为图像和视频做什么？arXiv预印本arXiv：2106.11297，2021。1、4[26] Andreas Steiner，Alexander Kolesnikov，Zhai Xiaohua，Ross Wightman，Jakob Uszkoreit，Lucas Beyer。如何训练你的vit？数据、增强和正则化。arXiv预印本arXiv：2106.10270，2021。4[27] 理查德·萨顿苦的教训，2019年3月。http：//www. 不完整的想法。net/IncIdeas/BitterLesson. html，1. 4[28] Yi Tay ， Mostafa Dehghani ， Vamsi Aribandi ， JaiGupta，Philip Pham，Zhen Qin，Dara Bahri，Da-ChengJuan，and Donald Metzler.OmniNet：来自transformers的全向表示。 arXiv预印本arXiv：2103.01075，2021。1、4[29] Yi Tay 、 Mostafa Dehghani 、 Jinfeng Rao 、 WilliamFedus 、 Samira Abnar 、 Hyung Won Chung 、 SharanNarang 、 Dani Yogatama 、 Ashish Vaswani 和 DonaldMetzler。有效扩展：来自预训练和微调转换器的见解arXiv预印本arXiv：2109.10686，2021。1[30] Ilya Tolstikhin，Neil Houlsby，Alexander Kolesnikov，Lu- cas Beyer ， Xiaohua Zhai ， Thomas Unterthiner ，Jessica Yung，Daniel Keysers，Jakob Uszkoreit，MarioLucic，et al. MLP- Mixer：An all-MLP architecture forvision.arXiv预印本arXiv：2105.01601，2021。1、4[31] Shen Yan ， Xuehan Xiong ， Anurag Arnab ， ZhichaoLu，Mi Zhang，Chen Sun，and Cordelia Schmid. 用于视 频 识 别 的 多 视 图 转 换 器 。 arXiv 预 印 本 arXiv ：2201.04288，2022。1