核心示例

核心示例托管在 GitHub Flax 仓库的 examples 目录中。

每个示例旨在 自包含且易于派生，同时在机器学习的不同领域再现相关结果。

如 #231 中所述，我们决定为所有示例采用标准模式，包括最简单的示例（如 MNIST）。这使得每个示例都更加详细，但是一旦你了解一个示例，你就了解了所有示例的结构。当您派生这些示例时，拥有单元测试和集成测试也非常有用。

下面的一些示例有一个链接“交互式🕹”，可以让你在 Colab 中直接运行它们。

图像分类

• MNIST - Interactive🕹: 用于 MNIST 分类卷积神经网络（具有简单代码）。

• ImageNet - Interactive🕹: 带有权重衰减的 ImageNet 上的 Resnet-50（具有多主机 SPMD、自定义预处理、检查点、动态缩放、混合精度）。

强化学习

• 近端策略优化: 学习玩 Atari 游戏（具有单主机 SPMD、RL 设置）。

自然语言处理

• 用于数字加法的序列到序列:（具有简单代码、LSTM 状态处理、即时数据生成）。

• 词性标注: 使用通用依存关系数据集的简单 Transformer 编码器模型。

• 情感分类: 使用 LSTM 模型。

• 在 WMT 上训练的 Transformer 编码器/解码器模型: 翻译英语/德语（具有多主机 SPMD、动态分桶、注意力缓存、打包序列、在 GCP 上进行 TPU 训练的方案）。

• 在十亿词基准上训练的 Transformer 编码器: 用于自回归语言建模，基于上述 WMT 示例。

生成模型

• 变分自动编码器: 在二值化 MNIST 上训练（具有简单代码、vmap）。

图模型

• 图神经网络: 在来自开放图基准的 ogbg-molpcba 上进行分子预测。

贡献核心 Flax 示例

大多数 GitHub 上的核心 Flax 示例 遵循 Flax 开发团队发现对 Flax 项目有效的结构。该团队努力使这些示例易于探索和派生。特别是（根据 GitHub 问题 #231）

• README: 包含指向论文、命令行、TensorBoard 指标的链接。

• 重点: 一个示例关于一个模型/数据集。

• 配置: 我们使用 ml_collections.ConfigDict 存储在 configs/ 下。

• 测试: 可执行的 main.py 加载 train.py，其中包含 train_test.py。

• 数据: 从 TensorFlow Datasets 读取。

• 独立: 每个目录都是自包含的。

• 要求: 版本固定在 requirements.txt 中。

• 样板: 通过使用 clu 减少。

• 交互式: 可以使用 Colab 探索该示例。