核心示例

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

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

如 #231 中讨论的那样，我们决定为所有示例（包括最简单的示例（如 MNIST））采用标准模式。这使得每个示例都稍微冗长一些，但是一旦您了解了一个示例，您就了解了所有示例的结构。当您 fork 这些示例时，进行单元测试和集成测试也非常有用。

以下某些示例具有“交互式🕹”链接，可让您直接在 Colab 中运行它们。

图像分类

• MNIST - 交互式🕹：用于 MNIST 分类的卷积神经网络（具有简单的代码）。

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

强化学习

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

自然语言处理

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

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

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

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

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

生成模型

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

图建模

• 图神经网络：来自 Open Graph Benchmark 的 ogbg-molpcba 上的分子预测。

为核心 Flax 示例做贡献

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

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

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

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

• 测试：可执行的 main.py 加载具有 train_test.py 的 train.py。

• 数据：从 TensorFlow 数据集读取。

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

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

• 样板代码：通过使用 clu 来减少。

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