Repolode
提交项目
google

google/flax

面向 JAX 的神经网络库，强调灵活性，提供 Flax NNX API 用于构建、检查和调试模型。

在 GitHub 查看官网
Star
7,366
Fork
849
未关闭 Issue
526
最近提交
前天

这是什么？

Flax 是构建在 JAX 之上的神经网络库及配套生态。其较新的 Flax NNX API 允许用普通 Python 对象定义模型，支持引用共享与可变状态，使网络的创建、检查、调试和分析更加直观。更早的 Flax Linen API 发布于 2020 年，是 NNX 的前身。

库中提供 Linear、Conv、BatchNorm、LayerNorm、GroupNorm、多头注意力、LSTM 和 GRU 单元、Dropout 等常用层，并包含多设备复制训练、序列化与检查点、指标计算以及数据预取到设备等工具。示例涵盖 MNIST 以及 Gemma 语言模型的推理。

Flax 的设计思路是让用户直接修改训练循环，而不是依赖框架额外添加的功能。可通过 PyPI 安装，需要 Python 3.8 及以上版本。

核心功能

  • 01采用 NNX API，用普通 Python 对象定义模型，支持引用共享与可变状态
  • 02内置 Linear、Conv、归一化层、注意力、循环单元和 Dropout 等常用层
  • 03提供多设备复制训练、检查点保存、指标计算和设备数据预取工具
  • 04附带 MNIST 和 Gemma 语言模型推理等示例
  • 05可通过 PyPI 安装，基于 JAX 运行

适用场景

  • 在 JAX 上研究和快速原型化新的神经网络结构
  • 为实验编写自定义训练循环，不受框架约束
  • 通过 JAX 在 CPU、GPU 或 TPU 上训练模型
  • 借助 MNIST 和 Gemma 示例学习基于 JAX 的深度学习

快速上手

  1. 1

    先按照官方 JAX 安装说明为你的硬件安装 JAX

  2. 2

    从 PyPI 安装 Flax

    pip install flax

使用前须知

  • Flax 依赖 JAX，因此你还需要为自己的硬件安装并配置 JAX。

标签

#jax#neural-network#deep-learning#machine-learning#python#research

常见问题

flax 是什么？
面向 JAX 的神经网络库，强调灵活性，提供 Flax NNX API 用于构建、检查和调试模型。 Flax 是构建在 JAX 之上的神经网络库及配套生态。其较新的 Flax NNX API 允许用普通 Python 对象定义模型，支持引用共享与可变状态，使网络的创建、检查、调试和分析更加直观。更早的 Flax Linen API 发布于 2020 年，是 NNX 的前身。
flax 可以免费使用吗？用的什么许可证？
flax 以 Apache-2.0 许可证开源。用于商业产品前，请先确认许可证条款。
flax 还在维护吗？
活跃维护中。flax 最近一次提交是在 2026年11月5日。
flax 适合用来做什么？
在 JAX 上研究和快速原型化新的神经网络结构 为实验编写自定义训练循环，不受框架约束 通过 JAX 在 CPU、GPU 或 TPU 上训练模型 借助 MNIST 和 Gemma 示例学习基于 JAX 的深度学习
flax 怎么安装和运行？
先按照官方 JAX 安装说明为你的硬件安装 JAX 从 PyPI 安装 Flax: pip install flax
在哪里获取 flax？
源代码和安装说明在 GitHub 上：https://github.com/google/flax

概要由 AI 根据项目 README 生成；GitHub 数据每日更新。

相关项目