deep-learning - 亚麻 (google) 和 dm-haiku (deepmind) 之间的主要区别是什么？

Question

flax之间的主要区别是什么和 dm-haiku ?
从他们的描述来看:

Flax，JAX 的神经网络库

Haiku，受 Sonnet 启发的 JAX 神经网络库

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

UPD .
找到了 explanation关于与 dm-haiku 开发者的区别:

Flax is a bit more batteries included, and comes with optimizers, mixed precision and some training loops (I am told these are decoupled and you can use as much or as little as you want). Haiku aims to just solve NN modules and state management, it leaves other parts of the problem to other libraries (e.g. optax for optimization).

Haiku is designed to be a port of Sonnet (a TF NN library) to JAX. So Haiku is a better choice if (like DeepMind) you have a significant amount of Sonnet+TF code that you might want to use in JAX and you want migrating that code (in either direction) to be as easy as possible.

I think otherwise it comes down to personal preference. Within Alphabet there are 100s of researchers using each library so I don't think you can go wrong either way. At DeepMind we have standardised on Haiku because it makes sense for us. I would suggest taking a look at the example code provided by both libraries and seeing which matches your preferences for structuring experiments. I think you'll find that moving code from one library to another is not very complicated if you change your mind in the future.

最初的问题仍然相关。

Answer 1

我最近遇到了同样的问题，我喜欢 Haiku，因为我认为它们的实现(参见 Flax Dense() 与 Haiku Linear() )更接近原始 JAX 精神(即链接 init 和 predict 函数并跟踪参数在 Pytrees 中)，这使我更容易修改内容。
但是，如果您不想深入修改内容，最好的选择方法是在 CNNs + LSTMs with Flax/Haiku 上找到一篇不错的博客文章并坚持下去。我的一般观点是，即使我更喜欢 Haiku(+ Optax + Rlax + Chex + ...)的构建方式，但两个库都非常接近。