使用 Flax 的代码仓库#
以下代码库使用 Flax,并提供了训练框架和大量示例。在许多情况下,您还可以找到预训练权重。
🤗 Hugging Face#
🤗 Hugging Face 是一个非常流行的用于构建、训练和部署最先进的机器学习模型的库。这些模型可以应用于文本、图像和音频。在组织了 JAX/Flax 社区周之后,他们现在在其存储库中拥有超过 5,000 个 Flax/JAX 模型。
🥑 DALLE Mini#
🥑 DALLE Mini 是一个基于 Transformer 的文本到图像模型,用 JAX/Flax 实现,遵循 OpenAI 原创 DALLE 论文中的思想。
Scenic#
Scenic 是一个用于计算机视觉研究及其他领域的代码库/库。Scenic 的主要重点是基于注意力的模型。Scenic 已成功用于开发多种模态的分类、分割和检测模型,包括图像、视频、音频以及它们的多种模态组合。
Big Vision#
Big Vision 是一个旨在用于使用 Cloud TPU VM 或 GPU 机器训练大规模视觉模型的代码库。它基于 Jax/Flax 库,并使用 tf.data 和 TensorFlow Datasets 来实现可扩展和可重复的输入管道。这是 ViT、MLP-Mixer、LiT、UViM 和许多其他模型的原始代码库。