flax (google) 和 dm-haiku (deepmind) 之间的主要区别是啥？

Posted 2023-02-16

【中文标题】flax (google) 和 dm-haiku (deepmind) 之间的主要区别是啥？

【英文标题】：What is the main difference between flax (google) and dm-haiku (deepmind)?flax (google) 和 dm-haiku (deepmind) 之间的主要区别是什么？

【发布时间】：2021-07-22 22:31:33

【问题描述】：

flax 和 dm-haiku 之间的主要区别是什么？

根据他们的描述：

Flax，用于 JAX 的神经网络库 Haiku，一个受 Sonnet 启发的 JAX 神经网络库

问题：

我应该选择哪个基于 jax 的库来实现，比如说 DeepSpeech 模型（由 CNN 层 + LSTM 层 + FC 组成）和 ctc-loss？

---

UPD。

发现explanation与dm-haiku开发者的区别：

Flax 包含更多的电池，并带有优化器、混合精度和一些训练循环（我听说这些是解耦的，您可以根据需要使用多少）。 Haiku 旨在解决 NN 模块和状态管理，它将问题的其他部分留给其他库（例如用于优化的 optax）。

Haiku 被设计为将 Sonnet（一个 TF NN 库）移植到 JAX。因此，如果（如 DeepMind）您有大量的 Sonnet+TF 代码，您可能希望在 JAX 中使用并且您希望尽可能轻松地迁移该代码（在任一方向上），那么 Haiku 是一个更好的选择。

我认为否则归结为个人喜好。在 Alphabet 中，每个图书馆都有 100 名研究人员使用，所以我认为无论哪种方式都不会出错。在 DeepMind，我们对 Haiku 进行了标准化，因为它对我们有意义。我建议查看两个库提供的示例代码，看看哪些符合您对结构化实验的偏好。我想你会发现，如果你以后改变主意，将代码从一个库移动到另一个库并不是很复杂。

---

原来的问题仍然相关。

【参考方案1】：

我最近遇到了同样的问题，我喜欢 Haiku，因为我认为它们的实现（参见 Flax Dense() 与 Haiku Linear() ）更接近原始 JAX 精神（即链接 init 和 predict 函数并保持跟踪Pytrees 中的参数），这使我更容易修改。

但是，如果您不想深入修改内容，最好的选择方法是找到一篇关于 CNNs + LSTMs with Flax/Haiku 的不错的博客文章并坚持下去。我的普遍看法是，即使我更喜欢 Haiku (+ Optax + Rlax + Chex + ...) 的模块化构建方式，这两个库也非常接近。

【讨论】：

感谢您的回答。 dm-haiku 的语法更接近 tensorflow 但 flax 类似于 pytorch 的说法，你同意吗？

JAX/FLAX/Haiku 在我看来比 Tensorflow/Keras 更容易。我不知道 pytorch，但无论如何我不会将它与 FLAX 进行比较。对我来说，第一个选择是 TensorFlow、PyTorch 或 JAX，然后，如果您选择 JAX，第二个选择是 Flax/Trax/Haiku/...（或 Keras/Theano...，如果您选择 TensorFlow）。不过，这可能是一个过时的观点！